[论文解读] Learning Equivariant Energy Based Models with Equivariant Stein Variational Gradient Descent

Priyank Jaini, Lars Holdijk+1 more|arXiv (Cornell University)|Jun 15, 2021

Adversarial Robustness in Machine Learning参考文献 36被引用 5

一句话总结

该论文提出了等变Stein变分梯度下降（E-SVGD），一种具备对称性感知能力的采样方法，通过使用等变核函数，能够高效地从对称性不变的概率密度中生成多样化且高质量的样本。此外，该研究进一步提出了基于E-SVGD训练的等变能量模型（E-EBMs），在分子生成、图像建模和粒子系统任务中，通过结合几何深度学习与群等变神经网络，显著提升了采样效率、泛化能力与收敛速度。

ABSTRACT

We focus on the problem of efficient sampling and learning of probability densities by incorporating symmetries in probabilistic models. We first introduce Equivariant Stein Variational Gradient Descent algorithm -- an equivariant sampling method based on Stein's identity for sampling from densities with symmetries. Equivariant SVGD explicitly incorporates symmetry information in a density through equivariant kernels which makes the resultant sampler efficient both in terms of sample complexity and the quality of generated samples. Subsequently, we define equivariant energy based models to model invariant densities that are learned using contrastive divergence. By utilizing our equivariant SVGD for training equivariant EBMs, we propose new ways of improving and scaling up training of energy based models. We apply these equivariant energy models for modelling joint densities in regression and classification tasks for image datasets, many-body particle systems and molecular structure generation.

研究动机与目标

• 开发一种显式结合对称性不变性的采样方法，以提升对称性不变概率密度的采样多样性与效率。
• 通过引入在对称群作用下保持结构一致的等变核函数，解决标准SVGD在采样对称分布时的局限性。
• 设计等变能量模型（E-EBMs），使能量函数与采样过程均尊重数据的底层对称性。
• 证明将对称性感知学习与对称性感知采样相结合，可在生成建模任务中实现更快收敛与更好泛化性能。
• 在分子结构生成、具有旋转对称性的图像数据集以及多体粒子系统等多样化领域验证该框架的有效性。

提出的方法

• 提出等变SVGD（E-SVGD），一种基于粒子的采样方法，利用等变核函数演化粒子，使其分布保持在群作用下的不变性。
• 基于Stein恒等式，在再生核希尔伯特空间（RKHS）中推导出核化梯度流，其中核函数被构造为在对称群作用下一致变换。
• 定义等变能量模型（E-EBMs），其能量函数由等变神经网络（如E-GNN、可导向CNN、群等变网络）参数化，以尊重数据的对称性。
• 通过对比发散法训练E-EBMs，利用E-SVGD生成负样本，确保采样过程与能量函数具有相同的对称性。
• 使用群表示（如SO(3)、C4、Sn）定义粒子与函数在对称操作下的变换方式，实现模型与推理过程中的一致等变性。
• 采用基于轨道的密度分解方法，对商空间上的分布进行建模，降低复杂度并提升采样效率。

实验结果

研究问题

• RQ1 我们能否设计一种采样算法，显式保持目标分布的对称性，从而实现更高效、更多样化的采样？
• RQ2 在能量函数与采样过程均引入等变性，对能量模型的训练稳定性与泛化能力有何影响？
• RQ3 相较于标准SVGD，使用E-SVGD作为训练采样器在对称生成建模中，能否显著提升收敛速度与样本质量？
• RQ4 具有等变架构的E-EBMs在分子结构生成与对称性约束下的图像生成任务中，是否能超越训练数据实现良好泛化？
• RQ5 在涉及对称数据的分类与生成任务中，对称性感知模型相较于非等变基线模型表现如何？

主要发现

• 在高维与长程相互作用场景下，E-SVGD在采样对称密度时，相比标准SVGD展现出更优的样本多样性与质量。
• 使用E-SVGD训练的等变能量模型收敛更快、泛化能力更强，尤其在数据稀缺或高度对称的设定下表现显著优于标准EBMs。
• 在QM9数据集上训练的E-EBMs生成的分子展现出正确的3D键合排列，包括芳香环与构造异构体，径向分布函数与真实数据分布高度一致。
• 在具有C4对称性的FashionMNIST数据集上，使用E-SVGD训练的等变EBM相比使用原始SVGD或非等变架构的模型，实现了更高的测试准确率与更快收敛速度。
• 将E-SVGD用作训练采样器，可实现E-EBMs在复杂对称领域（如多体粒子系统）中的稳定且高效的对比发散训练。
• 实证结果表明，使用E-SVGD训练的E-EBMs在所有评估基准中均生成更真实、结构更一致的样本，涵盖图像生成、分子设计与粒子系统建模。