[论文解读] Learning Interpolations between Boltzmann Densities

Bálint Máté, François Fleuret|arXiv (Cornell University)|Jan 18, 2023

Gaussian Processes and Bayesian Inference被引用 5

一句话总结

该论文提出了一种连续归一化流的新颖训练目标，通过时间依赖的向量场和能量函数插值，用玻尔兹曼分布之间的插值替代反向KL散度。该方法通过最小化连续性方程的违反程度，减少了模式崩溃，并提升了在高斯混合分布和量子力学双阱势等多模态目标上的训练效率。

ABSTRACT

We introduce a training objective for continuous normalizing flows that can be used in the absence of samples but in the presence of an energy function. Our method relies on either a prescribed or a learnt interpolation $f_t$ of energy functions between the target energy $f_1$ and the energy function of a generalized Gaussian $f_0(x) = ||x/σ||_p^p$. The interpolation of energy functions induces an interpolation of Boltzmann densities $p_t \propto e^{-f_t}$ and we aim to find a time-dependent vector field $V_t$ that transports samples along the family $p_t$ of densities. The condition of transporting samples along the family $p_t$ is equivalent to satisfying the continuity equation with $V_t$ and $p_t = Z_t^{-1}e^{-f_t}$. Consequently, we optimize $V_t$ and $f_t$ to satisfy this partial differential equation. We experimentally compare the proposed training objective to the reverse KL-divergence on Gaussian mixtures and on the Boltzmann density of a quantum mechanical particle in a double-well potential.

研究动机与目标

• 解决在使用反向KL散度训练归一化流以估计多模态目标分布下期望时出现的模式崩溃问题。
• 通过基于连续能量函数插值的物理启发式训练目标，克服反向KL散度的局限性。
• 实现在无需真实样本的情况下进行训练，仅依赖能量函数评估及其梯度。
• 通过优化满足学习到的能量插值路径上连续性方程的时间依赖向量场，同时提升样本质量和训练速度。
• 在具有挑战性的目标（如多模态高斯分布和量子力学玻尔兹曼分布）上验证该方法的有效性。

提出的方法

• 定义时间依赖的能量函数 $ f_t = (1-t)f_0 + t f_1 + t(1-t)f^\theta(t) $，其中 $ f_0 $ 为广义高斯分布，$ f_1 $ 为目标能量函数。
• 诱导一族玻尔兹曼分布 $ p_t \propto e^{-f_t} $，并学习一个时间依赖的向量场 $ V_t $，以沿该族分布传输样本。
• 将正确传输的条件表述为连续性方程 $ \partial_t p_t + \nabla \cdot (p_t V_t) = 0 $，该方程将 $ V_t $ 与 $ p_t $ 联系起来。
• 通过最小化连续性方程的逐点违反程度来优化 $ V_t $ 和 $ f_t $，使用损失函数 $ \mathcal{L} = \int \mathcal{E} + \mathcal{E}^2 \, dt $，其中 $ \mathcal{E} $ 为残差误差。
• 使用具有共享架构的神经网络（MLPs）参数化 $ f_t $ 和 $ V_t $，并利用自动微分计算梯度。
• 使用Adam优化器配合余弦学习率衰减进行训练，并通过50步的四阶龙格-库塔法积分流。

实验结果

研究问题

• RQ1 在仅能获取能量函数评估的情况下，基于连续性方程的训练目标是否能减少归一化流中的模式崩溃？
• RQ2 与反向KL散度最小化相比，该方法在KL散度、有效样本大小和模式覆盖率方面的性能如何？
• RQ3 学习非线性插值 $ f_t $ 在基底与目标能量函数之间是否优于线性插值？
• RQ4 该方法在高维多模态目标上的训练速度和样本质量方面扩展性如何？
• RQ5 该方法能否推广至复杂物理系统，例如双阱势中量子粒子的玻尔兹曼分布？

主要发现

• 在高斯混合分布和量子双阱目标上，该方法的KL散度显著低于反向KL训练。
• 在高斯混合基准上，该方法将有效样本大小提升了最多2.5倍。
• 模式覆盖率得到显著改善，流模型成功学习到二维高斯混合分布的全部四个模式，未出现模式崩溃。
• 训练收敛更快且更稳定，仅训练10,000轮的模型性能即优于反向KL训练250,000轮的模型。
• 使用学习到的非线性插值 $ f_t $ 显著优于线性插值，尤其在能量景观的高曲率区域表现更优。
• 该方法自然适配PINN框架，连续性误差可作为残差损失，表明其在PDE约束学习中具有更广泛的应用潜力。