Event Generation with Parallel Langevin Sampling and Learned Stein Diagnostics

This paper proposes a parallel underdamped Langevin sampling method enhanced by learned Stein discrepancy diagnostics and neural network surrogate initialization to efficiently generate unweighted high-multiplicity collider events while minimizing expensive matrix-element evaluations.

The Gist

Imagine you are trying to find the best seats in a massive, dark stadium (the "phase space") for a huge concert. The "best seats" are the ones where the crowd is thickest, but the stadium is so big and the crowd distribution is so complicated that you can't see the whole picture at once. In particle physics, this is like trying to simulate a collision between two protons to see what particles fly out. The goal is to generate a list of "events" (like snapshots of the collision) that perfectly match the laws of physics, without any bias.

The problem is that the "best seats" (the most likely outcomes) are often hidden in tiny, hard-to-reach corners of the stadium. Traditional methods are like throwing darts blindfolded: you throw thousands of darts, but most hit empty seats. You have to throw so many that it takes forever to get a few good ones. This is called "rejection sampling," and it becomes a nightmare when the concert has many more particles involved (high multiplicity).

This paper proposes a smarter way to find those seats using Parallel Langevin Sampling and a Learned Stein Diagnostic. Here is how it works, using simple analogies:

1. The Hiking Team (Parallel Langevin Chains)

Instead of one person throwing darts, imagine sending out hundreds of hikers (chains) at the same time.

• The Hikers: Each hiker carries a backpack with a map (the target density) and a compass (the gradient). They don't just walk randomly; they use the compass to feel the slope of the terrain. If the ground slopes down toward a "crowded" area, they walk that way.
• The Momentum: These hikers have "momentum." If they are walking downhill, they don't stop immediately when the slope flattens; they keep gliding. This helps them cross small valleys and hills quickly without getting stuck.
• The Parallel Strategy: The researchers run thousands of these hikers simultaneously on powerful computers (GPUs). Crucially, they don't wait for a hiker to wander around for a long time and get confused (which creates "autocorrelation"). Instead, they let each hiker take a specific number of steps and then stop. They keep only the final position of each hiker and discard the path they took. This ensures every "event" they collect is fresh and independent.

2. The Coach with a Stopwatch (Learned Stein Diagnostic)

The big question is: How long should the hikers walk before we stop them?

• If they stop too early, they are still stuck near the starting line and haven't found the good seats.
• If they stop too late, they wasted time walking in circles.

The paper introduces a "Coach" (the Learned Stein Diagnostic). This Coach is an AI that watches the hikers. It doesn't know the exact map of the stadium, but it can compare the hikers' current positions to the "ideal" distribution of where they should be.

• The Coach uses a special test (the Stein Discrepancy) to measure how far off the hikers are from the perfect distribution.
• When the Coach sees that the hikers have finally settled into the right pattern (the discrepancy drops to near zero), it blows the whistle. This tells the system exactly how many steps were needed to reach "relaxation" (the point where the samples are valid).

3. The Training Wheels (Neural Network Surrogates)

Even with the Coach, the hikers still have to climb the mountain using the real, heavy map, which is slow and expensive to compute.

• The Shortcut: The researchers trained a simple AI "Surrogate" (a training wheel) on a small, cheap set of data. This AI learns to guess the shape of the stadium very quickly, though not perfectly.
• The Strategy: The hikers start their journey using this cheap, fast AI map. They get a head start and move quickly toward the right area. Once they are close, the researchers switch them to the real, heavy map for just a few final steps to get the perfect answer.
• The Result: This "warm start" drastically reduces the number of expensive, real calculations needed. It's like letting a student practice on a simulator before taking the real driving test.

What Did They Find?

The team tested this method on a specific particle collision process ($u\bar{u} \to Z + ng$), which involves a Z boson and varying numbers of gluons (particles).

• Efficiency: They found that the hikers only needed a modest number of steps to reach the "good seats," even as the number of particles increased.
• Accuracy: The final list of events they generated matched the results from traditional, trusted methods (called Vegas) perfectly.
• Speed: Using the "Training Wheels" (surrogate) cut the number of expensive calculations by a huge margin (e.g., for 3 gluons, they went from needing 115 steps to just 55).

The Bottom Line

This paper shows that instead of blindly throwing darts, we can send out a team of guided hikers, use an AI coach to tell us exactly when they are ready, and give them a head start with a cheap simulator. This makes generating complex particle physics events much faster and more efficient, which is crucial for the future of high-energy physics experiments like the Large Hadron Collider.

Technical Summary

Technical Summary: Event Generation with Parallel Langevin Sampling and Learned Stein Diagnostics

Problem Statement
Precision collider phenomenology faces a significant computational bottleneck in generating Monte Carlo (MC) event samples, particularly for high-multiplicity final states. The dominant cost arises from sampling the hard proton-proton collision differential cross section, where matrix-element evaluations are expensive. Current standard methods, largely based on the Vegas algorithm with physics-inspired improvements, rely on rejection sampling. This approach suffers from low unweighting efficiencies as multiplicity increases, often degrading by several orders of magnitude. While recent machine learning approaches (e.g., normalizing flows) have improved efficiency, they remain fundamentally dependent on rejection sampling and its associated inefficiencies.

Methodology
The authors propose an alternative framework based on parallelized Markov Chain Monte Carlo (MCMC) using the Underdamped Langevin Algorithm (ULA), augmented by a data-driven convergence diagnostic.

• Parallel Underdamped Langevin Sampling: Instead of running a single long chain, the method evolves many independent Markov chains in parallel on hardware accelerators. To eliminate within-chain autocorrelation, only the terminal state of each chain is retained as an unweighted event. The chains utilize an underdamped Langevin dynamic, retaining momentum between steps and applying a Metropolis-Hastings (MH) correction after each proposal to correct for discretization bias.
• Mirror Space Dynamics: Since the physical phase space is bounded ($x \in [0, 1]^d$) while ULA is naturally formulated on $\mathbb{R}^d$, the evolution is performed in "mirror space" using an element-wise transform (inverse error function). This preserves the convergence properties of the unconstrained algorithm.
• Learned Stein Discrepancy (LSD) as a Diagnostic: To determine when the ensemble of chains has reached stationarity (i.e., the target distribution), the authors employ a Learned Stein Discrepancy. Unlike kernel-based Stein discrepancies which scale quadratically with sample size ($n^2$), the LSD uses a neural network critic to estimate the discrepancy linearly with the number of samples. This allows for efficient monitoring of convergence.
• Handling Sharp Boundaries: Physical cross sections involve sharp phase-space cuts. To ensure the Stein operator remains valid in the presence of these hard boundaries, a scalar boundary function $h(x)$ is introduced. This function vanishes on the cut surface, ensuring the boundary contribution to the Stein operator is handled correctly.
• Surrogate Initialization: To reduce the number of expensive exact matrix-element and gradient evaluations, a neural network surrogate is trained on a limited set of phase-space points. Chains are first evolved using this surrogate to accelerate equilibration, then switched to the exact target density for a shorter sequence of steps before termination.

Key Contributions

1. Algorithmic Framework: The integration of parallel ULA chains with a learned Stein discrepancy provides a method for generating strictly unweighted events without relying on rejection sampling.
2. Convergence Diagnostic: The application of the Learned Stein Discrepancy as a practical, data-driven estimator for relaxation time ($t_{rel}$), allowing the system to determine when to stop evolving chains and accept samples.
3. Surrogate Acceleration: Demonstrating that simple neural network surrogates can substantially reduce the number of exact-target evaluations required for relaxation, acting as a "warm start" for the MCMC chains.
4. Boundary Handling: A specific formulation of the Stein operator that accounts for sharp phase-space cuts common in collider physics.

Results
The method was applied to tree-level $u\bar{u} \to Z + ng$ event generation for gluon multiplicities $n = 0, 1, 2, 3$.

• Convergence: The Learned Stein Discrepancy successfully tracked the convergence of the chains. For the studied multiplicities, relaxation required only a modest number of exact-target ULA steps, with only mild growth as multiplicity increased.
• Efficiency Gains: The use of surrogate initialization significantly reduced the required number of exact matrix-element evaluations. For example, for $Z+3g$, the relaxation time dropped from 115 steps (without surrogate) to 55 steps (with surrogate).
• Validation: The terminal samples from the ULA chains reproduced the distributions of physical observables (e.g., transverse momentum $p_T$, jet resolution scale $\sqrt{d_{01}}$, and rapidity) within statistical uncertainties when compared against independent Vegas estimates.

Significance and Claims
The paper claims that gradient-based parallel MCMC is a practically viable alternative to rejection-based event generation, particularly in regimes where unweighting efficiency degrades severely with multiplicity. The authors emphasize that their method generates strictly unweighted events, avoiding the residual weights often found in rejection sampling when maximum weights are capped to maintain efficiency.

The computational overhead of the method is described as negligible compared to the cost of matrix-element evaluations. The approach is naturally parallel, making it well-suited for GPU/TPU execution. The authors note that while the current study focuses on a fixed partonic channel, tree-level matrix elements, and a simple phase-space map, the combination of parallel chains, explicit relaxation-time diagnostics, and surrogate initialization offers a concrete route toward event generation that minimizes evaluations of the expensive exact target. The authors remain modest, stating that broader comparisons with current methods and scaling behavior must ultimately be tested within full event generator programs.