Name
Abstract | ||
|---|---|---|
We present a general approach to batching arbitrary computations for accelerators such as GPUs. We show orders-of-magnitude speedups using our method on the No U-Turn Sampler (NUTS), a workhorse algorithm in Bayesian statistics. The central challenge of batching NUTS and other Markov chain Monte Carlo algorithms is data-dependent control flow and recursion. We overcome this by mechanically transforming a single-example implementation into a form that explicitly tracks the current program point for each batch member, and only steps forward those in the same place. We present two different batching algorithms: a simpler, previously published one that inherits recursion from the host Python, and a more complex, novel one that implemenents recursion directly and can batch across it. We implement these batching methods as a general program transformation on Python source. Both the batching system and the NUTS implementation presented here are available as part of the popular TensorFlow Probability software package. |
Year | Venue | DocType |
|---|---|---|
2020 | MLSys | Conference |
Citations | PageRank | References |
0 | 0.34 | 0 |
Authors | ||
5 | ||
Name | Order | Citations | PageRank |
|---|---|---|---|
| Alexey Radul | 1 | 35 | 8.90 |
| Brian Patton | 2 | 9 | 2.08 |
| Dougal Maclaurin | 3 | 255 | 9.76 |
| Matthew D. Hoffman | 4 | 1117 | 60.21 |
| Rif Saurous | 5 | 148 | 10.49 |