PaperlayerView exact revision on arXiv

Source-linked AI summary

Learning the Geometry of Collider Events with Metric-Aware Deep Sets

Lauren Hay, Rishabh Jain, Matt LeBlanc, Jennifer Roloff

arXiv:2609.12024v1hep-phcs.LG

TL;DR

Exact optimal transport is costly for large pairwise analyses, and learned surrogates need not preserve distance geometry. The paper introduces a Deep Sets EMD surrogate with selected geometric constraints. It reports percent-level errors, higher throughput than exact CPU solvers in large workloads, and fewer triangle-inequality violations.

  • Problem

    Exact OT is costly in large pairwise analyses, while low regression error alone does not ensure that learned distances preserve geometric properties.

    Source: §1 · p. 1

  • Method

    The Metric-Aware Particle Flow Network is a Deep Sets surrogate for collider-event EMD that enforces non-negativity, exchange symmetry, and zero self-distance.

    Source: §1 · p. 2

  • Results

    Median relative errors are of order 1% across both collision processes and all three event-generation stages, with higher measured throughput than exact CPU OT implementations in large pairwise workloads.

    Source: §4 · p. 9

  • Takeaways & Limitations

    Across 106 held-out event triplets, architectural constraints reduced observed triangle-inequality violations, indicating improved geometric fidelity beyond the properties directly enforced.

    Source: §4 · p. 9

    The surrogate does not guarantee separation of distinct events, the triangle inequality, or IRC safety, and relative errors increase at small distances.

    Source: §4 · p. 9

Abstract

from arXiv · show

Optimal transport gives structured data a geometry, but exact evaluation is costly in large pairwise analyses that exploit relationships among distances. Learned surrogates are faster, but need not preserve this metric structure. We develop a Deep Sets surrogate for OT between variable-size weighted point clouds that enforces non-negativity, exchange symmetry, and zero self-distance, leaving the triangle inequality unconstrained. Applied to the Energy Mover's Distance between collider events in a particle physics application, the Metric-Aware Particle Flow Network achieves percent-level mean absolute percentage error while significantly improving inference throughput over other exact and approximate methods surveyed. The architectural constraints are found to improve properties that are not explicitly enforced: across $10^6$ held-out event triplets, triangle-inequality violations fall from 199 for a matched unconstrained network to 2, and the maximum from 149.5 to 5.8 GeV. These results demonstrate that targeted inductive biases can yield fast neural surrogates with substantially improved geometric fidelity.

1 Introduction

Exact optimal transport is costly in large pairwise analyses, while learned costs can be numerically accurate yet geometrically inconsistent. The paper introduces a constrained Deep Sets surrogate for collider-event EMD and evaluates both its accuracy and geometric fidelity.

  • Exact OT becomes costly in large pairwise analyses because each of O(n^2) comparisons requires a transport problem whose cost grows with object support size.These pairwise relationships underpin nearest-neighbor searches, clustering, and embedding.Source: §1 · p. 1
  • Numerical accuracy does not ensure that predicted distances are non-negative, symmetric, zero on the diagonal, or triangle-consistent, motivating architectural constraints that target only part of the geometry.The study asks how much target geometry such constraints can restore without enforcing the full metric structure.Source: §1 · p. 1
  • Prior fast approaches include entropic and sliced OT approximations, linearized transport representations, learned metric embeddings, and neural EMD approximations.This work instead directly supervises a pairwise regressor while encoding part of the target geometry in a permutation-invariant architecture.Source: §1 · p. 1
  • Collider events provide a demanding test because each is a variable-size weighted point cloud, and typical exact comparisons optimize over roughly a 200×200 particle-level cost matrix.The same mathematical setting also covers OT between other weighted point clouds and empirical distributions.Source: §1 · pp. 1–2
  • 199 to 2 triangle-inequality violations and 149.5 to 5.8 GeV maximum violations accompany sub-percent mean absolute percentage error and a growing dense-pair throughput advantage over exact CPU solvers.These results are reported for 106 held-out triplets.Source: §1 · p. 2

2 Methodology

The methodology defines collider events and EMD targets, examines infrared and collinear safety, and constructs and trains a metric-aware Deep Sets surrogate.

2.1 Data and target distance

  • The target is an unbalanced OT metric that transports transverse momentum between variable-size event point clouds and penalizes unmatched weight, yielding EMD values in GeV.The angular ground distance uses β = 1 and R = Rmax = 11.64.Source: §2.1 · p. 2
  • Z+jets and t¯t samples span hard-scatter, parton-shower, and hadronization stages; training uses Z+jets PS pairs, while other stages and t¯t test multiplicity generalization.Events, rather than pairs, are partitioned across splits so no event is shared.Source: §2.1 · p. 2

2.2 Infrared and Collinear Safety

  • IRC safety requires insensitivity to vanishing-momentum particles and coincident collinear splittings; the MA-PFN constraints do not enforce these limits, so the study tests limiting and finite-perturbation responses.EMD proximity implies closeness for IRC-safe observables.Source: §2.2 · p. 3

2.3 Metric-Aware Particle Flow Network Implementation

  • The model uses a shared particle encoder and is trained on 6,395,676 exact EMD labels with a hybrid relative- and absolute-error objective.The selected checkpoint is epoch 648, chosen by validation objective.Source: §2.3 · pp. 3–4
  • The MA-PFN forms an exact latent pseudometric from shared event embeddings and multiplies it by a positive, exchange-invariant learned gauge, leaving triangle inequality unconstrained.A joint head processes symmetric and antisymmetric event representations; the baseline instead encodes the union of both events as one set.Source: §2.3 · pp. 3–4

3 Results

The MA-PFN maintains strong held-out EMD prediction accuracy while improving several learned distance properties and scaling favorably against exact and approximate alternatives. Its infrared and collinear responses track the exact EMD to smaller perturbations than the baseline, but plateau rather than enforcing those limits.

3.1 Predictive Accuracy

  • 0.890% MAPE on 798,216 held-out Z+jets PS pairs, versus 1.124% for the matched baseline PFN; MA-PFN lowers each reported error by approximately 20%.MA-PFN also achieves 0.679 GeV MAE and 0.943 GeV RMSE, versus 0.843 GeV and 1.176 GeV for the baseline PFN.Source: §3.1 · p. 4

3.2 Learned Metric Properties

  • 2 of 106 MA-PFN triplets violate the triangle inequality, versus 199 for the baseline PFN; the largest positive residual falls from 149.5 GeV to 5.8 GeV.Triangle consistency is learned empirically, not guaranteed by the architecture.Source: §3.2 · p. 5, Figure 2 · p. 6

3.3 Learned Infrared and Collinear Safety

  • For soft emission and collinear splitting, MA-PFN follows the exact response to smaller perturbations than the baseline before reaching a nonzero floor.The large-angle peak near Δϕ = 2π is caused by azimuthal periodicity; only the Δϕ → 0 region tests collinear safety, and the closer tracking also appears for uneven splittings.Source: §3.3 · pp. 6–7, Figure 3 · p. 7

3.4 Timing Performance

  • MA-PFN throughput advantage over the fastest exact implementation grows from 1.5× at median multiplicity 4 to 1.8 × 10^3 at median multiplicity 202.Mixed-precision MA-PFN takes 0.038–0.057 s on one NVIDIA A40, or 17–26 million pairs per second; the reported advantage is not hardware-normalized.Source: §3.4 · p. 7, Figure 4 · p. 8

3.4.1 Comparison of approximate methods

  • The MA-PFN encodes each of 2,000 unique events once, then evaluates pairwise heads on cached 64-dimensional latents, making the pair stage independent of particle multiplicity.This amortizes multiplicity-dependent processing into event encoding and explains the nearly flat timing curve.Source: §3.4.1 · p. 8, Figure 5 · p. 9
  • MA-PFN reduces median absolute relative error to approximately 0.85% while processing at least 1.4 × 10^7 pairs/s, Pareto-dominating the tested Sinkhorn configuration.Both MA-PFN precision modes also have lower median absolute relative error and greater aggregate throughput than the tested SWD and Sinkhorn configurations.Source: §3.4.1 · p. 8, Figure 5 · p. 9

4 Conclusion

The MA-PFN achieves roughly 1% median relative error across the tested collision processes and generation stages, with higher measured throughput than exact CPU solvers in large pairwise workloads. Its constraints improve observed triangle-inequality behavior, but it remains an approximate EMD surrogate with geometric and small-distance limitations.

4 Conclusion

  • Median relative errors are of order 1% across both collision processes and all three event-generation stages, with higher measured throughput than exact CPU optimal-transport implementations in large pairwise workloads.The throughput advantage increases rapidly with event multiplicity.Source: §4 · p. 9
  • On 106 held-out event triplets, MA-PFN reduces observed triangle-inequality violation rate by ∼100× versus baseline PFN and lowers the largest violation from 149.5 to 5.8 GeV.In paired-event infrared and 50/50 collinear studies, it has lower median signed-response error than PFN at every sampled perturbation scale. Both models approach nonzero plateaus in infrared and collinear limits rather than the exact EMD response.Source: §4 · p. 9
  • Separation of distinct events, the triangle inequality, and IRC safety are not guaranteed, and relative errors increase in the small-distance regime.The MA-PFN is an approximate surrogate rather than an exact EMD replacement.Source: §4 · p. 9

A.1 Hardware and execution settings

The benchmark uses one Oscar-cluster CPU node for exact solvers and one NVIDIA A40 for GPU methods. Timings include method-specific preparation and output production, with GPU batch sizes selected separately for each method and sample.

A.1 Hardware and execution settings

  • Exact solvers run in Float64 on one Oscar-cluster node with 48 physical CPU cores, while GPU methods run sequentially on one NVIDIA A40.The tested A40 devices are multiple physical GPUs of the same model, so the comparison reflects tested deployments rather than kernels on one identical device.Source: §A.1 · p. 13
  • Timing includes method-specific preparation, distance evaluation, synchronization, and an ordered host-resident output, but excludes disk input, imports, solver construction, worker startup, and JIT warm-up.MA-PFN is benchmarked in FP32 and automatic mixed precision; Figure 4b reports FP32 and Figure 5b includes both modes.Source: §A.1 · p. 13
  • GPU batch sizes are tuned independently by method and sample, rejecting candidates above 80% device memory and selecting the smallest batch that reaches at least 98% of peak median throughput.MA-PFN pair-stage memory is nearly independent of particle multiplicity because each pair uses two fixed-width cached latents, unlike SWD and Sinkhorn.Source: §A.1 · p. 13

A.2 Repetitions and numerical exceptions

The timing protocol uses repeated trials, with cache handling differing between cache-aware GPU methods and MA-PFN. Two exact-solver exceptions affect reported points in the largest workloads.

A.2 Repetitions and numerical exceptions

  • The main 106-pair workload uses three to seven CPU and SWD repetitions, three Sinkhorn repetitions, and five MA-PFN repetitions for each precision and sample.Additional CPU and SWD repetitions are run up to seven unless cumulative measured time reaches 30 seconds.Source: §A.2 · p. 13
  • For cache-aware GPU methods, the first repetition builds and transfers reusable per-event state, while later repetitions reuse it; MA-PFN resident-cache medians are reported and cold-cache preparation times recorded separately.Figures report median wall time, and throughput is pair count divided by wall time.Source: §A.2 · p. 13
  • The default Wasserstein C++ solver failed on one t¯t PS million-pair workload pair, and EnergyFlow.jl ns64 points are omitted for the two largest workloads after one non-optimal result each.The plotted Wasserstein value uses a recovery configuration with a tenfold higher small-cost tolerance that passed exact-solver validation.Source: §A.2 · p. 13

B Additional Details about Multiset Deep Sets Theorem

The Deep Sets theorem provides a representation for permutation-invariant functions of sets, and the multiset extension accommodates observables over multiple events. For the two-event EMD, a transformed input basis supports exchange symmetry and the construction’s other imposed properties.

B Additional Details about Multiset Deep Sets Theorem

  • The standard Deep Sets theorem states that a permutation-invariant function on a bounded set can be approximated using a learned representation.The stated conditions include a bounded set of elements in Rd and a bounded interval as the function’s output space.Source: B Additional Details about Multiset… · pp. 13–14
  • The multiset extension treats permutations within each set as invariant, supporting functions of several independent sets.The passage presents this as the more natural general extension for observables over multiple sets.Source: B Additional Details about Multiset… · p. 14
  • The set representation zX is formed with learned continuous encoder and decoder functions, and the multiset theorem describes observables defined over multiple events, including EMD between two events.The passage identifies Φ as the encoder and F as the decoder.Source: B Additional Details about Multiset… · p. 14
  • The input transformation separates symmetric and antisymmetric terms; averaging and rescaling the resulting outputs enforce exchange symmetry, positivity, and identity of indiscernibles.The transformation is invertible, so the passage states it incurs no loss of expressivity.Source: B Additional Details about Multiset… · p. 14

Related papers

Loading 2609.12024v1…