# Parallelizing MCMC Across the Sequence Length
This repository accompanies the paper "Parallelizing MCMC Across the Sequence Length" submitted to NeurIPS 2025.

**A note on data:** for the parallel MALA experiments, we provided a whitened version of the German Credit Problem dataset in the `mala/data` directory. This is to maximize reproducibility and avoid numerical discrepancies from different interpretations of whitening and centering, etc.

**A note on dependencies:** please see `requirements.txt` for a full list of packages and their corresponding versions. Please note that `JAX` is particularly sensitive to cluster environments and CUDA-related packages/settings. 

**A note on compute:** all experiments for this paper were run on a Linux-based SLURM cluster using compute nodes with either single NVIDIA A100 or H100 GPUs with 80 GB of RAM. As described in the paper, all timing experiments were run using a single H100 with 80 GB of RAM unless otherwise specified.

**Helper files:** for the parallel MALA experiments, the main helper function is the `seq1d` function located in `src/algs/deer.py`, which runs the quasi-DEER and DEER algorithms given an MCMC update function. We gratefully acknowledge that our code was based on the original DEER codebase by Lim et al. (2024), available at [this GitHub repository](https://github.com/machine-discovery/deer), which accompanies their paper ["Parallelizing Non-linear Sequential Models over the Sequence Length"](https://arxiv.org/abs/2309.12252) (ICLR 2024).

The `mala` directory contains all code for running the parallel MALA experiments and generating all main text and supplementary figures:
- The `data` folder contains `german-credit_d=24_whiten=True_camera-ready.csv`, a whitened version of the German Credit Problem dataset (also available from [UCI]([http://archive.%20ics.%20uci.%20edu/ml](https://archive.ics.uci.edu/datasets)) and [TensorFlow Datasets](https://www.tensorflow.org/datasets)).
- For full reproducibility, the `logs` folder contains all log files for all of the parallel MALA experiments, which are used when generating parallel MALA main text and supplementary figures.
- The `src` folder contains the `deer.py` and other related files for implementing the DEER and quasi-DEER algorithms.
- The `truth` folder contains our set of *gold standard* reference samples from the posterior distribution of the Bayesian logistic regression model fitted to the German Credit Problem. These were generated using TensorFlow Probability's state-of-the-art NUTS sampler with dual-averaging stepsize adaptation. We use these sample as the ground truth when computing Maximum Mean Discrepancy (MMD).
- `0. Creating Data Structures.ipynb` creates starter directories and data structures to facilitate logging and saving of experimental results.
- `1. Generating Ground Truth Samples.ipynb` generates the *gold standard* reference samples described above.
- `2. Selecting Sigma for MMD.ipynb` uses [Gretton et al. (2012)](https://jmlr.csail.mit.edu/papers/volume13/gretton12a/gretton12a.pdf)'s recommendation to calibrate the MMD kernel hyperparameter $\sigma$ based on the median heuristic for the Bayesian logistic regression problem on the German Credit dataset.
- `3. Analyzer.ipynb` produces summaries of our raw experimental results that are then used for making parallel MALA main text and supplementary figures.
- `4. Main Text Panels.ipynb` contains everything for reproducing our main text parallel MALA figures.
- `5. Supplementary Figures.ipynb` contains everything for reproducing our supplementary materials parallel MALA figures.
- The two files `runner.py` and `runner.sh` run all fully-converged parallel MALA experimental settings with memory-efficient quasi-DEER.
- The two files `runner_ablation.py` and `runner_ablation.sh` run all fully-converged parallel MALA experimental settings with naive quasi-DEER (i.e., without the memory-efficient optimizations).
- The two files `runner_fulltrace.py` and `runner_fulltrace.sh` run all parallel MALA experimental settings with memory-efficient quasi-DEER with `full_trace=True` so we have complete samples from each Newton iteration, and not just at convergence.

**Additional Experiments:** Scripts for running the parallel and sequential Gibbs sampler, HMC with sequential or paralleal leapfrog integration, and parallel HMC across the sequence are found in the `scripts` folder. 

**License:** all assets in our paper and accompanying repository are under CC BY-NC 4.0.
