Coupling normalizing flows allow for fast sampling and density evaluation, making them the tool of choice for probabilistic modeling of physical systems. However, the standard coupling architecture precludes endowing flows that operate on the Cartesian coordinates of atoms with the SE(3) and permutation invariances of physical systems. This work proposes a coupling flow that preserves SE(3) and permutation equivariance by performing coordinate splits along additional augmented dimensions. At each layer, the flow maps atoms' positions into learned SE(3) invariant bases, where we apply standard flow transformations, such as monotonic rational-quadratic splines, before returning to the original basis. Crucially, our flow preserves fast sampling and density evaluation, and may be used to produce unbiased estimates of expectations with respect to the target distribution via importance sampling. When trained on the DW4, LJ13 and QM9-positional datasets, our flow is competitive with equivariant continuous normalizing flows, while allowing sampling two orders of magnitude faster. Moreover, to the best of our knowledge, we are the first to learn the full Boltzmann distribution of alanine dipeptide by only modeling the Cartesian positions of its atoms. Lastly, we demonstrate that our flow can be trained to approximately sample from the Boltzmann distribution of the DW4 and LJ13 particle systems using only their energy functions.

通过在附加的扩充维度上执行坐标分割，本文提出了一种保持SE(3)和排列等变性的耦合流，能够快速采样和密度评估，并可用于产生对目标分布的期望的无偏估计。同时，本文显示该流模型在相同应用上比等变连续归一化流快两个数量级，并首次成功学习仅通过对原子的笛卡尔位置进行建模从而获得丙氨酸二肽的玻尔兹曼分布。最后，本文展示了该流模型可以通过仅使用粒子系统的能量函数在DW4和LJ13粒子系统的玻尔兹曼分布上进行近似采样。

SE(3)等变增强耦合流