Automatic Differentiation via Effects and Handlers in OCaml
Machine learning, artificial intelligence, and scientific modelling have driven the demand for tools that enable derivative based optimization. Automatic differentiation (AD) is a family of algorithms used to calculate the derivatives of programs with only a constant factor slowdown. Effects and handlers are a powerful programming language control flow construct based on delimited continuations. Mainstream programming languages are increasingly incorporating effects and handlers, notably OCaml 5.0.
We show that effects and handlers are a great match for implementing AD algorithms while maintaining asymptotic efficiency. In particular, effects and handlers allow for succinctness in the presence of the intrinsic complex control flow of AD. We implement four AD algorithms in OCaml 5.0 using effects and handlers. We provide benchmarks to empirically show that we can reach the correct asymptotic complexity for forward and reverse mode AD. We also provide a real-world comparison by adding our implementation to a preexisting benchmark suite which includes systems TensorFlow and PyTorch, and show that our implementation is competitive with systems based on comparable methods, second only to PyTorch and Tensorflow.
Slides (ML 2024 slides.pdf) | 1.67MiB |
Fri 6 SepDisplayed time zone: Amsterdam, Berlin, Bern, Rome, Stockholm, Vienna change
11:00 - 12:30 | |||
11:00 30mTalk | Is there a use for linear types? ML | ||
11:30 30mTalk | Automatic Differentiation via Effects and Handlers in OCaml ML Jesse Sigal University of Edinburgh Pre-print File Attached | ||
12:00 30mTalk | Rethinking the Value Restriction ML Stephen Dolan Jane Street |