Efficiently Vectorized MCMC on Modern Accelerators
Authors: Hugh Dance, Pierre Glaser, Peter Orbanz, Ryan P Adams
ICML 2025 | Venue PDF | Archive PDF | Plain Text | LLM Run Details
| Reproducibility Variable | Result | LLM Response |
|---|---|---|
| Research Type | Experimental | We implement several popular MCMC algorithms as FSMs, including Elliptical Slice Sampling, HMC-NUTS, and Delayed Rejection, demonstrating speed-ups of up to an order of magnitude in experiments. |
| Researcher Affiliation | Academia | 1Gatsby Unit, University College London, UK 2Department of Computer Science, Princeton University, USA. Correspondence to: Hugh Dance <EMAIL>. |
| Pseudocode | Yes | Algorithm 1 MCMC algorithm with sample function ... Algorithm 2 step function for FSM ... Algorithm 3 FSM MCMC algorithm ... Algorithm 4 bundled step for FSM with S1, S2 ... Algorithm 5 amortized step for FSM with function g |
| Open Source Code | Yes | Code can be found at https://github.com/hwdance/jax-fsm-mcmc. |
| Open Datasets | Yes | Gaussian Process Regression on the Real Estate Dataset (Yeh, 2018)... UCI Machine Learning Repository. DOI: https://doi.org/10.24432/C5J30W. |
| Dataset Splits | No | The paper does not provide specific dataset split information (exact percentages, sample counts, citations to predefined splits, or detailed splitting methodology) needed to reproduce the data partitioning. |
| Hardware Specification | Yes | All experiments are run in JAX on an NVIDIA A100 GPU with 32GB CPU memory. |
| Software Dependencies | No | The paper mentions software like JAX, Num Pyro, Black JAX, TensorFlow, PyTorch, and Flax but does not provide specific version numbers for these components that would be necessary for reproduction. |
| Experiment Setup | Yes | We use a N(x, 0.1) proposal distribution with M = 100 tries per sample and draw 10,000 samples per chain. ... Normal priors σ, τ, λ N(0, 1), (so the ellipse is drawn using N(0, I)). ... We use a pre-tuned step-size with acceptance rate 0.85 and identity mass matrix. ... average results over 128 chains of 1000 samples, with hyperparameters pre-tuned over 400 warm-up steps. |