Differentiable Matrix Elements with MadJax
This provides a step towards differentiable programming in high energy physics, facilitating optimization pipelines for researchers in that domain.
The authors tackled the challenge of incorporating high energy physics simulation software into gradient-based learning by developing MadJax, a tool that generates differentiable matrix elements for scattering processes, enabling applications like simulation-based inference and normalizing flow modeling.
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.