JetFlow: Multi-Domain Applications

• JetFlow is a multifaceted term defining flow-based methods that span collider jet generation, speculative decoding in language models, and jet-actuated fluid dynamics control.
• In particle physics, JetFlow employs mass-conditioned normalizing flows to generate high-dimensional jet point clouds, significantly improving invariant mass reproduction and generation fidelity.
• Extensions of JetFlow include speculative decoding frameworks for LLMs and flow-based surrogates in CFD, yielding substantial speedups and enhanced performance in high-energy physics and fluid control.

JetFlow is not a single fixed designation in current arXiv usage. The supplied literature shows that it names, or explicitly motivates, several distinct but structurally related research programs. In particle physics, the term most directly denotes a likelihood-based normalizing-flow generator for collider jets that conditions on invariant jet mass and optionally on constituent multiplicity to reproduce high-level jet correlations on JetNet (Käch et al., 2022). Later high-energy-physics work extends that “JetFlow” idea to constituent-level flow matching on JetClass, many-jet matrix-element sampling, multimodal continuous–discrete particle-cloud generation, and flow-based surrogates for jet-induced hydrodynamic response in heavy-ion collisions (Birk et al., 2023, Bothmann et al., 23 Jun 2025, Faroughy et al., 1 Sep 2025, Wu et al., 17 May 2026). Outside collider phenomenology, JetFlow also names a speculative decoding framework for LLMs, where it denotes a causal parallel tree drafter for head-based speculative decoding (Hu et al., 16 Jun 2026). In the supplied fluid-dynamics materials, by contrast, “JetFlow” functions mainly as an application perspective for jet-actuated CFD control rather than as the formal title of the underlying software framework (Xiao et al., 1 Aug 2025, Wang et al., 2024).

1. Terminological scope

The term spans several technical domains, but the usages are not interchangeable. In particle physics, JetFlow primarily refers to flow-based models for jets as particle clouds or as many-body phase-space objects. In language-model systems, it refers to an inference-time acceleration method. In fluid mechanics, it appears as a jet-flow-control framing for HPC-coupled reinforcement learning and differentiable CFD rather than as a single canonical package.

Usage of “JetFlow”	Domain	Core object
“JetFlow: Generating Jets with Conditioned and Mass Constrained Normalising Flows” (Käch et al., 2022)	Collider-jet generation	Mass-conditioned RQS normalizing flow on JetNet
“Flow Matching Beyond Kinematics: Generating Jets with Particle-ID and Trajectory Displacement Information” (Birk et al., 2023)	Particle-cloud generation	JetFlow-style EPiC-FM on JetClass
“Efficient many-jet event generation with Flow Matching” (Bothmann et al., 23 Jun 2025)	Matrix-element event generation	JetFlow-type phase-space sampler
“Multimodal Generative Flows for LHC Jets” (Faroughy et al., 1 Sep 2025)	Multimodal jet generation	Continuous–discrete flow for particle clouds
“A flow-matching generative model for event-by-event jet-induced hydro response in high-energy heavy-ion collisions” (Wu et al., 17 May 2026)	Heavy-ion phenomenology	Conditional flow for jet-induced hydro spectra
“JetFlow: Breaking the Scaling Ceiling of Speculative Decoding with Parallel Tree Drafting” (Hu et al., 16 Jun 2026)	LLM inference	Causal parallel speculative drafter

Within the collider-physics literature, TopicFlow explicitly places Käch et al.’s JetFlow in the broader landscape of “using normalizing flows to model jets,” while distinguishing it from topic disentangling of quark and gluon mixtures from mixed samples (Dolan et al., 2022).

2. Mass-conditioned normalizing flows for collider jets

The 2022 JetFlow of Käch et al. addresses fast generation of collider jets as high-dimensional point clouds with a tractable likelihood. It is trained and evaluated on the JetNet datasets, where each sample is a jet of about $1$ TeV energy with up to $30$ constituents and three relative kinematic features per constituent: $\eta_i^{\mathrm{rel}}$, $\phi_i^{\mathrm{rel}}$, and $p_{T,i}^{\mathrm{rel}}$. This yields a $90$-dimensional representation per jet. A central global observable is the relative invariant mass,

$$(m^{\mathrm{rel}})^2 = \left(\sum_i E_i^{\mathrm{rel}}\right)^2 - \left(\sum_i \vec p_i^{\mathrm{rel}}\right)^2 = \frac{m_{\text{jet}}^2}{p_{T,\text{jet}}^2},$$

whose distribution is highly sensitive to constituent-level correlations and therefore serves as a stringent test of generative fidelity (Käch et al., 2022).

Architecturally, JetFlow is a coupling-layer normalizing flow with rational–quadratic monotonic spline transforms rather than affine couplings. The overall map $f$ is invertible, and the density is computed by the change-of-variables formula,

$$p_X(x) = p_Z(f(x))\left|\det \frac{\partial f(x)}{\partial x}\right|.$$

The model uses coupling layers in the RealNVP/Glow style, but replaces affine transforms by spline transforms parameterized by widths, heights, and derivatives. Hyperparameters scanned include the number of bins $K=4\ldots 7$, the bound $30$0, the number of coupling layers from $30$1 up to $30$2, and $30$3 residual blocks in the conditioner networks. The conditioners are fully connected feed-forward networks with residual blocks, ReLU activation, dropout, and batch normalization.

The central empirical finding is that an unconditioned “Vanilla NF” models inclusive per-particle marginals reasonably well but fails to reproduce the invariant mass distribution. JetFlow remedies this by conditioning the spline-parameter networks on the invariant mass $30$4, and optionally on the constituent count $30$5. In practice, the conditioner sees $30$6 rather than only $30$7. Variable-length jets are embedded into a fixed $30$8-particle representation by zero-padding, with noise of order $30$9 added to padded entries to avoid degeneracies incompatible with bijectivity. When conditioning on multiplicity, particles beyond the requested $\eta_i^{\mathrm{rel}}$0 are forced to zero after sampling. Since the model does not generate $\eta_i^{\mathrm{rel}}$1 or $\eta_i^{\mathrm{rel}}$2 by itself, mass sampling is handled by an empirical CDF interpolated with monotone piecewise cubic Hermite polynomials, and $\eta_i^{\mathrm{rel}}$3 is sampled from its empirical histogram.

JetFlow further studies a mass-consistency penalty,

$\eta_i^{\mathrm{rel}}$4

added to the negative log-likelihood to form $\eta_i^{\mathrm{rel}}$5. This yields three model classes: VNF, CNF, and CCNF. On JetNet, the unconditioned VNF shows $\eta_i^{\mathrm{rel}}$6 of order $\eta_i^{\mathrm{rel}}$7–$\eta_i^{\mathrm{rel}}$8, whereas mass conditioning reduces that discrepancy by almost an order of magnitude. For gluon jets, the reported numbers are $\eta_i^{\mathrm{rel}}$9 for VNF, about $\phi_i^{\mathrm{rel}}$0 for CNF(m), and $\phi_i^{\mathrm{rel}}$1 for CCNF(m). Training converges within $\phi_i^{\mathrm{rel}}$2–$\phi_i^{\mathrm{rel}}$3 hours on an NVIDIA P100 GPU, and sampling is reported at $\phi_i^{\mathrm{rel}}$4 per jet for batch size $\phi_i^{\mathrm{rel}}$5. The model is therefore positioned as a fast, stable, likelihood-based alternative whose main conceptual contribution is to make a key Lorentz-invariant summary statistic explicit in the generative mechanism rather than leaving it to be learned implicitly.

3. Constituent-level, multimodal, and particle-cloud generalizations

Several later papers explicitly describe their constructions as JetFlow-style extensions of flow-based jet generation. The first generative model trained on JetClass uses a permutation-equivariant continuous normalizing flow trained with flow matching, conditioned on jet type, and extends generation beyond kinematics to include particle-ID information, charge, and track displacement variables (Birk et al., 2023). Each jet is represented by up to $\phi_i^{\mathrm{rel}}$6 constituents with $\phi_i^{\mathrm{rel}}$7 features per constituent: three relative kinematic coordinates, four impact-parameter-related quantities, charge, and five particle-ID flags. The architecture scales the EPiC-FM design to $\phi_i^{\mathrm{rel}}$8 EPiC layers with hidden dimension $\phi_i^{\mathrm{rel}}$9 and global dimension $p_{T,i}^{\mathrm{rel}}$0. It is trained on $p_{T,i}^{\mathrm{rel}}$1 million jets, with about $p_{T,i}^{\mathrm{rel}}$2 million trainable parameters, AdamW at learning rate $p_{T,i}^{\mathrm{rel}}$3, batch size $p_{T,i}^{\mathrm{rel}}$4, $p_{T,i}^{\mathrm{rel}}$5 epochs, and midpoint ODE integration with $p_{T,i}^{\mathrm{rel}}$6 function evaluations at generation time. The model reproduces per-particle kinematics, particle-type fractions, and several jet-level observables well, but the paper identifies hadronic-top $p_{T,i}^{\mathrm{rel}}$7 and impact-parameter tails as particularly difficult; in the classifier-based evaluation, hadronic tops are the most distinguishable generated class.

A subsequent multimodal flow makes the discrete component explicit instead of treating one-hot particle identities as merely continuous surrogates (Faroughy et al., 1 Sep 2025). In that construction, each particle has continuous kinematics $p_{T,i}^{\mathrm{rel}}$8 and a discrete flavor token

$p_{T,i}^{\mathrm{rel}}$9

with jets padded and masked to $90$0 particles. The continuous part is trained by flow matching, while the discrete part is generated by a continuous-time Markov jump bridge derived from a multivariate random telegraph process. The network is a “Multimodal ParticleFormer” with mode-specific branches for kinematics and flavor, followed by a fused transformer block; it has about $90$1 million parameters and is trained on $90$2 million jets with $90$3 for validation, using $90$4 A100 GPUs. Relative to an EPiC-FM baseline, this multimodal model improves jet mass, $90$5, $90$6, $90$7, $90$8, and jet charge, while underperforming on very rare leptonic particle classes. The technical significance is that the continuous–discrete coupling is no longer delegated to post-hoc argmax projection but is embedded into the generative path itself.

4. Phase-space transport and jet-induced medium response

The JetFlow label is also used conceptually for flow-based samplers and surrogates that act on jet-related objects other than constituent clouds. In many-jet event generation, Flow Matching is applied to the random numbers used by the Pepper partonic event generator to sample momenta and helicities for high-multiplicity Drell–Yan and $90$9 final states (Bothmann et al., 23 Jun 2025). The phase-space dimension is

$$(m^{\mathrm{rel}})^2 = \left(\sum_i E_i^{\mathrm{rel}}\right)^2 - \left(\sum_i \vec p_i^{\mathrm{rel}}\right)^2 = \frac{m_{\text{jet}}^2}{p_{T,\text{jet}}^2},$$0

reaching $$(m^{\mathrm{rel}})^2 = \left(\sum_i E_i^{\mathrm{rel}}\right)^2 - \left(\sum_i \vec p_i^{\mathrm{rel}}\right)^2 = \frac{m_{\text{jet}}^2}{p_{T,\text{jet}}^2},$$1 in the studied channels. The learned model is a Continuous Normalizing Flow trained with Flow Matching, optionally conditioned on helicity labels via embeddings. Its role is not to replace the matrix element, but to remap the sampling density so that the resulting weights are flatter and unweighting becomes more efficient. For the highest multiplicities studied, the reported improvements relative to Vegas are large: for Drell–Yan plus five jets, $$(m^{\mathrm{rel}})^2 = \left(\sum_i E_i^{\mathrm{rel}}\right)^2 - \left(\sum_i \vec p_i^{\mathrm{rel}}\right)^2 = \frac{m_{\text{jet}}^2}{p_{T,\text{jet}}^2},$$2 rises to $$(m^{\mathrm{rel}})^2 = \left(\sum_i E_i^{\mathrm{rel}}\right)^2 - \left(\sum_i \vec p_i^{\mathrm{rel}}\right)^2 = \frac{m_{\text{jet}}^2}{p_{T,\text{jet}}^2},$$3 versus $$(m^{\mathrm{rel}})^2 = \left(\sum_i E_i^{\mathrm{rel}}\right)^2 - \left(\sum_i \vec p_i^{\mathrm{rel}}\right)^2 = \frac{m_{\text{jet}}^2}{p_{T,\text{jet}}^2},$$4, and for $$(m^{\mathrm{rel}})^2 = \left(\sum_i E_i^{\mathrm{rel}}\right)^2 - \left(\sum_i \vec p_i^{\mathrm{rel}}\right)^2 = \frac{m_{\text{jet}}^2}{p_{T,\text{jet}}^2},$$5 plus four jets it reaches $$(m^{\mathrm{rel}})^2 = \left(\sum_i E_i^{\mathrm{rel}}\right)^2 - \left(\sum_i \vec p_i^{\mathrm{rel}}\right)^2 = \frac{m_{\text{jet}}^2}{p_{T,\text{jet}}^2},$$6 versus $$(m^{\mathrm{rel}})^2 = \left(\sum_i E_i^{\mathrm{rel}}\right)^2 - \left(\sum_i \vec p_i^{\mathrm{rel}}\right)^2 = \frac{m_{\text{jet}}^2}{p_{T,\text{jet}}^2},$$7. The abstract summarizes these gains as factors of $$(m^{\mathrm{rel}})^2 = \left(\sum_i E_i^{\mathrm{rel}}\right)^2 - \left(\sum_i \vec p_i^{\mathrm{rel}}\right)^2 = \frac{m_{\text{jet}}^2}{p_{T,\text{jet}}^2},$$8 and $$(m^{\mathrm{rel}})^2 = \left(\sum_i E_i^{\mathrm{rel}}\right)^2 - \left(\sum_i \vec p_i^{\mathrm{rel}}\right)^2 = \frac{m_{\text{jet}}^2}{p_{T,\text{jet}}^2},$$9, respectively. In this usage, “JetFlow” refers to a flow-based transport map over many-jet phase space rather than to constituent-level jet synthesis.

A distinct heavy-ion application uses Flow Matching as a conditional surrogate for event-by-event jet-induced hydrodynamic response in Pb+Pb collisions (Wu et al., 17 May 2026). There the target is not a particle cloud but the final-state hadron spectrum

$f$0

generated by the jet-induced hydro response in CoLBT-hydro. The condition vector has dimension $f$1 and contains photon momentum, production position, and up to three reconstructed jets. The response spectrum is PCA-compressed to $f$2 latent coefficients, normalized to $f$3, and generated by a residual MLP velocity field trained with the linear interpolation

$f$4

Sampling uses a second-order Runge–Kutta integrator with $f$5 uniform steps. The surrogate reproduces the front wake and diffusion wake of the Mach-cone-like response, including the rapidity-asymmetry observable designed to cancel the hydro background. In the $f$6 GeV/c fit region, the effective temperatures are reported as $f$7 GeV and $f$8 GeV. The stated computational acceleration is approximately six orders of magnitude relative to full CoLBT-hydro. A plausible implication is that “JetFlow” can denote not only generative models of jets themselves, but also fast conditional surrogates for jet-induced medium effects.

5. Speculative decoding under the JetFlow name

In 2026, JetFlow becomes the title of a speculative decoding framework for autoregressive LLMs (Hu et al., 16 Jun 2026). Here the problem is not jet physics but the scaling ceiling of speculative decoding. If $f$9 is the draft window length, $$p_X(x) = p_Z(f(x))\left|\det \frac{\partial f(x)}{\partial x}\right|.$$0 the per-token acceptance rate, and $$p_X(x) = p_Z(f(x))\left|\det \frac{\partial f(x)}{\partial x}\right|.$$1 the relative per-token draft cost, the paper uses the stylized speedup expression

$$p_X(x) = p_Z(f(x))\left|\det \frac{\partial f(x)}{\partial x}\right|.$$2

This makes the core difficulty explicit: increasing draft budget helps only if acceptance remains high and drafting overhead remains low. Prior head-based methods face what the paper calls a causality–efficiency dilemma. Autoregressive drafters are path-conditioned but sequential; bidirectional or block-diffusion drafters are cheap but branch-agnostic and therefore internally inconsistent at tree scale.

JetFlow resolves this by attaching a lightweight causal parallel draft head to a frozen target model and imposing a tree-causal attention mask so that each node conditions on prefix tokens and its own ancestors, but not on siblings or descendants. For Qwen3-8B, the draft head fuses hidden states from target layers $$p_X(x) = p_Z(f(x))\left|\det \frac{\partial f(x)}{\partial x}\right|.$$3, projects them back to hidden size $$p_X(x) = p_Z(f(x))\left|\det \frac{\partial f(x)}{\partial x}\right|.$$4, and uses a $$p_X(x) = p_Z(f(x))\left|\det \frac{\partial f(x)}{\partial x}\right|.$$5-layer decoder with $$p_X(x) = p_Z(f(x))\left|\det \frac{\partial f(x)}{\partial x}\right|.$$6 attention heads, $$p_X(x) = p_Z(f(x))\left|\det \frac{\partial f(x)}{\partial x}\right|.$$7 KV heads, head dimension $$p_X(x) = p_Z(f(x))\left|\det \frac{\partial f(x)}{\partial x}\right|.$$8, and MLP size $$p_X(x) = p_Z(f(x))\left|\det \frac{\partial f(x)}{\partial x}\right|.$$9. All nodes across all depths are still processed in one forward pass. The resulting branchwise factorization is designed to align with the target model’s autoregressive factorization, unlike surrogate tree scores built from branch-agnostic marginals.

Training is performed by distillation on target-aligned sequences, preferably regenerated continuations from the target model, using a forward-KL objective rather than reverse-KL. The paper reports that reverse-KL causes a $K=4\ldots 7$0–$K=4\ldots 7$1 relative drop in speedup in ablations, whereas forward-KL and standard SFT are similar, with forward-KL slightly better. Empirically, the model reaches up to $K=4\ldots 7$2 speedup on MATH-500 and $K=4\ldots 7$3 on open-ended conversational workloads on H100 GPUs. At tree budget $K=4\ldots 7$4 on Qwen3-8B, MATH-500 reaches an accepted length $K=4\ldots 7$5. With vLLM integration, batch-size-$K=4\ldots 7$6 throughput on MATH-500 rises from $K=4\ldots 7$7 tokens/s for autoregressive decoding to $K=4\ldots 7$8 tokens/s for JetFlow, i.e. $K=4\ldots 7$9 speedup. The framework thus uses the JetFlow name to denote a causal block-parallel transport over speculative token trees rather than a generative flow in the normalizing-flow sense.

6. JetFlow as a fluid-dynamics perspective

In the supplied fluid-dynamics materials, “JetFlow” is used as an application perspective rather than as the formal title of the main framework. SmartFlow is the named system: a CFD-solver-agnostic deep reinforcement learning framework for HPC environments, built on Relexi, SmartSOD2D, SmartSim, and SmartRedis-MPI, and explicitly motivated by jet-based flow control problems such as synthetic jets, wall blowing/suction, and jet arrays on bluff bodies (Xiao et al., 1 Aug 2025). The architecture couples MPI-parallel CFD solvers to Python DRL algorithms by in-memory communication routines such as put_state, get_action, and put_reward, with a Gym-like CFDEnv on the Python side. In the 2D cylinder demonstration, two zero-net-mass-flux synthetic jets are mounted at $30$00 and $30$01, the control variable is the bounded scalar $30$02 with $30$03, observations are $30$04 pressure probes, and PPO uses MLP actor and critic networks with $30$05 hidden layers of $30$06 neurons and Tanh activations. Training uses $30$07 parallel FLEXI environments, $30$08 episodes, and about $30$09 hours on $30$10 AMD EPYC 7543 CPU cores. In the 3D SOD2D case with $30$11 jet pairs, the multi-agent setup achieves about $30$12 mean drag reduction and about $30$13 reduction of lift fluctuations.

A second line of work uses “JetFlow” as the name of an envisioned differentiable CFD workflow built around JAX-Fluids rather than as the title of the published software itself (Wang et al., 2024). JAX-Fluids is a Python, JAX-based solver for compressible single- and two-phase flow that solves the full compressible Navier–Stokes equations with a high-order Godunov-type finite-volume formulation, WENO/TENO-type reconstructions, HLLC Riemann solvers, and high-order central finite differences for viscous terms. Because mesh handling, fluxes, and time integration are expressed in JAX primitives, the whole solver is differentiable and compatible with jax.grad, jax.jit, and jax.pmap. On a HAWK-AI node with NVIDIA A100 GPUs, the paper reports wall-clock time below $30$14 s per time step on $30$15 GPUs for a $30$16-million-cell compressible turbulent channel-flow DNS in single precision, with performance comparable to other high-order codes such as FLEXI and GALÆXI. The same paper also introduces a differentiable disturbance-wall boundary condition for wall-normal blowing/suction, implemented without eval(), so that boundary parameters remain inside the AD graph. This suggests a version of “JetFlow” in which jet actuation, wall forcing, or porous-wall parameters are optimized end-to-end through a differentiable PDE solver rather than through black-box reinforcement learning.

The fluid-dynamics usage therefore differs from the particle-physics and LLM usages in one important respect: JetFlow is not a stable software name, but a descriptive shorthand for jet-actuated, flow-coupled learning systems on HPC platforms.