Staging Automatic Differentiation with Fusion
summary
Claims a compiler for reverse‑mode AD that is correct by construction and yields concrete speedups up to 110x speedup over an earlier interpreter approach, built by staging with Template Haskell staging and applying fusion laws plus monadic sharing/let insertion to avoid expression swell. It generates Haskell code, includes mechanized proofs of fusion, shows benchmarked wins against AD libraries (sometimes beating JAX, sometimes not), and reports occasional compile‑time OOMs.
AI-written summary · about
chapters (15)
- 0:00Staging automatic differentiation
- 1:26Fast and correct AD goals
- 3:55Reverse mode AD trade-offs
- 5:43Dual numbers and forward AD
- 7:34Fusion theorem for reverse AD
- 8:36Staging reverse AD with Template Haskell
- 11:04Evaluator as a compiler
- 13:15Compiler improvements with fusion
- 14:40Avoiding expression swell
- 16:11Optimizations via homomorphisms
- 17:14Constant folding extension
- 18:41Final compiler proof
- 19:23Benchmarking vs JAX and AD
- 21:54Summary and future work
- 23:42Q&A
related talks
01h 21mNon-local compiler transformations in the presence of dynamic dispatch
16mCompiling Symbolic Execution with Staging and Algebraic Effects
36mProgramming Languages for Automatic Differentiation: What Now?
26mAdding AD to Scheme by Differentiating the Interpreter
55mDeforestation and Program Fusion
24mCorrectness Meets Performance: From Agda to Futhark
29mReverse Automatic Differentiation for Accelerate (Extended Abstract)
29mParallelism-preserving automatic differentiation for second-order array languages