[Paper Review] 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 Applications4 citations

TL;DR

This paper introduces NEEMo, a differentiable geometric fitting framework that uses neural estimation of the Energy Mover’s Distance (EMD) via the Kantorovich-Rubinstein duality to enable exact, gradient-based optimization of EMD without ε-approximations. By leveraging a 1-Lipschitz neural network architecture, NEEMo enables precise, differentiable EMD computation for applications in high-energy particle physics, such as jet substructure clustering and event geometry fitting, achieving state-of-the-art alignment with target distributions in synthetic experiments.

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.

Motivation & Objective

• To develop a differentiable, exact method for computing the Energy Mover’s Distance (EMD) in high-energy particle physics, avoiding the ε-approximations of prior methods.
• To enable gradient-based optimization over the exact EMD by leveraging the Kantorovich-Rubinstein duality and 1-Lipschitz neural networks.
• To apply the framework to geometric fitting tasks such as clustering particle events into jets and aligning parametric shapes to target event distributions.
• To provide a flexible, differentiable framework for constructing precision geometric observables sensitive to substructure and detector effects.
• To explore the potential of NEEMo in next-generation collider experiments, such as the Electron-Ion Collider, where traditional clustering methods are suboptimal.

Proposed method

• The method uses the Kantorovich-Rubinstein dual formulation of the Wasserstein-1 metric to cast EMD computation as an optimization over 1-Lipschitz functions.
• A fully connected neural network with weight-norm constraints is used to parameterize the 1-Lipschitz function, ensuring exact Lipschitz bounds via operator norm constraints.
• The network is trained using gradient descent to maximize the difference in expectations of the function over the source and target distributions, yielding the exact EMD.
• The Kantorovich potential, derived from the trained network, induces forces on geometric parameters (e.g., circle centers), driving them toward alignment with the target distribution.
• The framework supports end-to-end, differentiable optimization of both the EMD and the geometric parameters of the source distribution.
• Small modifications to the KR dual formulation allow reliable and accurate EMD estimation without ε-Sinkhorn approximations.

Experimental results

Research questions

• RQ1 Can the Energy Mover’s Distance be computed exactly and differentiably using neural networks with exact 1-Lipschitz constraints?
• RQ2 How does the NEEMo framework improve upon ε-Sinkhorn approximations in EMD estimation for geometric fitting tasks?
• RQ3 To what extent can NEEMo accurately recover the true geometric structure of particle-collider events, such as jet substructure or circular/elliptical particle distributions?
• RQ4 Can NEEMo be used to perform unsupervised, event-by-event jet geometry fitting in a way that outperforms traditional clustering algorithms?
• RQ5 What is the potential impact of NEEMo on precision physics at future colliders like the Electron-Ion Collider, where standard methods are suboptimal?

Key findings

• NEEMo successfully fits three randomly initialized circles to a synthetic event with three circular particle clusters, achieving perfect alignment through gradient descent on the Kantorovich potential.
• The framework accurately fits parametric shapes (e.g., triangle and ellipse) to target distributions, demonstrating robustness to complex geometric structures.
• In N-subjet clustering experiments, NEEMo correctly identifies the true number of subjets (3, 4, or 5) by minimizing the EMD between fitted and true distributions.
• The EMD of the N-subjet fit was consistently lowest for the correct number of subjets, confirming the method’s sensitivity to true substructure.
• The framework enables differentiable optimization over the exact EMD, eliminating the need for ε-approximations used in prior methods like SHAPER’s Sinkhorn estimation.
• The approach shows strong potential for precision physics applications, including detecting percent-level fluctuations in event modifications due to hadronization or detector effects.