# Differentiable Matrix Elements with MadJax

arXiv: https://arxiv.org/abs/2203.00057 · HTML: https://mlsanity.com/paper/2203.00057/

## Key facts

- Cited by 8 papers in the map (data through 14 Oct 2026).
- Builds on 5 of its 26 references that are in the map.
- First posted on arXiv 28 Feb 2022.

## Builds on (most cited first, 5 of 5)

- [The frontier of simulation-based inference](https://mlsanity.com/paper/1911.01429/) (2019)
- [Machine Learning of coarse-grained Molecular Dynamics Force Fields](https://mlsanity.com/paper/1812.01736/) (2018)
- [Mining gold from implicit models to improve likelihood-free inference](https://mlsanity.com/paper/1805.12244/) (2018)
- [TorchMD: A deep learning framework for molecular simulations](https://mlsanity.com/paper/2012.12106/) (2020)
- [PDFFlow: parton distribution functions on GPU](https://mlsanity.com/paper/2009.06635/) (2020)

## Built on by (most cited first, 8 of 8)

- [Finetuning Foundation Models for Joint Analysis Optimization](https://mlsanity.com/paper/2401.13536/) (2024)
- [New directions for surrogate models and differentiable programming for High Energy Physics detector simulation](https://mlsanity.com/paper/2203.08806/) (2022)
- [Branches of a Tree: Taking Derivatives of Programs with Discrete and Branching Randomness in High Energy Physics](https://mlsanity.com/paper/2308.16680/) (2023)
- [Constraining the Higgs Potential with Neural Simulation-based Inference for Di-Higgs Production](https://mlsanity.com/paper/2405.15847/) (2024)
- [Differentiable Vertex Fitting for Jet Flavour Tagging](https://mlsanity.com/paper/2310.12804/) (2023)
- [Supercharging Simulation-Based Inference for Bayesian Optimal Experimental Design](https://mlsanity.com/paper/2602.06900/) (2026)
- [Flow Annealed Importance Sampling Bootstrap meets Differentiable Particle Physics](https://mlsanity.com/paper/2411.16234/) (2024)
- [Event Generation with Parallel Langevin Sampling and Learned Stein Diagnostics](https://mlsanity.com/paper/2606.14854/) (2026)

## Abstract (arXiv)

MadJax is a tool for generating and evaluating differentiable matrix elements of high energy scattering processes. As such, it is a step towards a differentiable programming paradigm in high energy physics that facilitates the incorporation of high energy physics domain knowledge, encoded in simulation software, into gradient based learning and optimization pipelines. MadJax comprises two components: (a) a plugin to the general purpose matrix element generator MadGraph that integrates matrix element and phase space sampling code with the JAX differentiable programming framework, and (b) a standalone wrapping API for accessing the matrix element code and its gradients, which are computed with automatic differentiation. The MadJax implementation and example applications of simulation based inference and normalizing flow based matrix element modeling, with capabilities enabled uniquely with differentiable matrix elements, are presented.