Poster
in
Workshop: NeurIPS 2024 Workshop: Machine Learning and the Physical Sciences
Synax: A Differentiable and GPU-accelerated Synchrotron Simulation Package
Kangning Diao · Zequn Li · Richard Grumitt · Yi Mao
Abstract:
We introduce Synax, a novel library for automatically differentiable simulations of Galactic synchrotron emission. Synax uses JAX's automatic differentiation (AD) mechanism, enabling precise computation of derivatives with respect to any model parameters. This feature facilitates powerful inference algorithms, such as Hamiltonian Monte Carlo (HMC) and gradient-based optimization, which enables inference over models that would otherwise be computationally prohibitive. Notably, we show that GPU acceleration brings a twenty-fold enhancement in efficiency, while HMC achieves a two-fold improvement over standard random walk Metropolis-Hastings when performing inference over a four-parameter test model.
Chat is not available.