[论文解读] Finding NEEMo: Geometric Fitting using Neural Estimation of the Energy Mover's Distance

O. Kitouni, N. S. Nolte+1 more|arXiv (Cornell University)|Sep 30, 2022

Medical Imaging Techniques and Applications被引用 4

一句话总结

该论文提出了 NEEMo，一种可微分的几何拟合框架，通过利用 Kantorovich-Rubinstein 对偶性对能量最优传输距离（EMD）进行神经估计，实现无需 ε-近似 的精确、基于梯度的 EMD 优化。通过采用 1-Lipschitz 神经网络架构，NEEMo 实现了高能粒子物理应用中精确且可微分的 EMD 计算，例如喷注子结构聚类和事件几何拟合，在合成实验中实现了与目标分布的最先进对齐效果。

ABSTRACT

A novel neural architecture was recently developed that enforces an exact upper bound on the Lipschitz constant of the model by constraining the norm of its weights in a minimal way, resulting in higher expressiveness compared to other techniques. We present a new and interesting direction for this architecture: estimation of the Wasserstein metric (Earth Mover's Distance) in optimal transport by employing the Kantorovich-Rubinstein duality to enable its use in geometric fitting applications. Specifically, we focus on the field of high-energy particle physics, where it has been shown that a metric for the space of particle-collider events can be defined based on the Wasserstein metric, referred to as the Energy Mover's Distance (EMD). This metrization has the potential to revolutionize data-driven collider phenomenology. The work presented here represents a major step towards realizing this goal by providing a differentiable way of directly calculating the EMD. We show how the flexibility that our approach enables can be used to develop novel clustering algorithms.

研究动机与目标

• 开发一种在高能粒子物理中精确计算能量最优传输距离（EMD）的可微分方法，避免先前方法中使用的 ε-近似。
• 通过利用 Kantorovich-Rubinstein 对偶性与 1-Lipschitz 神经网络，实现对精确 EMD 的基于梯度的优化。
• 将该框架应用于几何拟合任务，例如将粒子事件聚类为喷注，或将参数化形状与目标事件分布对齐。
• 提供一种灵活且可微分的框架，用于构建对子结构和探测器效应敏感的高精度几何可观测量。
• 探索 NEEMo 在下一代对撞机实验（如电子-离子对撞机）中的潜力，其中传统聚类方法表现不佳。

提出的方法

• 该方法利用 Wasserstein-1 距离的 Kantorovich-Rubinstein 对偶公式，将 EMD 计算转化为对 1-Lipschitz 函数的优化问题。
• 使用带权重范数约束的全连接神经网络来参数化 1-Lipschitz 函数，通过算子范数约束确保精确的利普希茨界。
• 通过梯度下降训练网络，以最大化函数在源分布与目标分布上的期望差异，从而获得精确的 EMD。
• 从训练网络中提取的 Kantorovich 潜势会在几何参数（如圆心）上施加作用力，推动其与目标分布对齐。
• 该框架支持对 EMD 与源分布几何参数的端到端、可微分优化。
• 对 KR 对偶公式的微小修改可实现无需 ε-Sinkhorn 近似的确保可靠且精确的 EMD 估计。

实验结果

研究问题

• RQ1 能否使用具有精确 1-Lipschitz 约束的神经网络，实现 EMD 的精确且可微分计算？
• RQ2 NEEMo 框架在几何拟合任务中对 EMD 估计的 ε-Sinkhorn 近似相比有何改进？
• RQ3 NEEMo 在多大程度上能准确恢复粒子对撞机事件的真实几何结构，如喷注子结构或圆形/椭圆形粒子分布？
• RQ4 NEEMo 是否可用于无监督的、逐事件的喷注几何拟合，且性能优于传统聚类算法？
• RQ5 NEEMo 对未来对撞机（如电子-离子对撞机）中的精密物理研究可能产生何种影响，尤其在标准方法表现不佳的场景下？

主要发现

• NEEMo 通过 Kantorovich 潜势的梯度下降，成功将三个随机初始化的圆拟合到具有三个圆形粒子簇的合成事件中，实现了完美对齐。
• 该框架能准确地将参数化形状（如三角形和椭圆）拟合到目标分布，展示了对复杂几何结构的鲁棒性。
• 在 N-子喷注聚类实验中，NEEMo 通过最小化拟合分布与真实分布之间的 EMD，正确识别出子喷注的真实数量（3、4 或 5 个）。
• 对于正确的子喷注数量，N-子喷注拟合的 EMD 值始终最低，证实了该方法对真实子结构的敏感性。
• 该框架实现了对精确 EMD 的可微分优化，消除了先前方法（如 SHAPER 的 Sinkhorn 估计）中使用的 ε-近似。
• 该方法在精密物理应用中展现出巨大潜力，包括检测由强子化或探测器效应引起的事件修改中的百分之一量级波动。