arXiv · 2212.09801
Denotationally Correct, Purely Functional, Efficient Reverse-mode Automatic Differentiation
Abstract
Reverse-mode differentiation is used for optimization, but it introduces references, which break the purity of the underlying programs, making them notoriously harder to optimize. We present a reverse-mode differentiation on a purely functional language with array operations. It is the first one to deliver a provably efficient, purely functional, and denotationally correct reverse-mode differentiation. We show that our transformation is semantically correct and verifies the cheap gradient principle. Inspired by PROPs and compilation to categories, we introduce a novel intermediate representation that we call 'unary form'. Our reverse-mode transformation is factored as a compilation scheme through this intermediate representation. We obtain provably efficient gradients by performing general partial evaluation optimizations after our reverse-mode transformation, as opposed to manually derived ones. For simple first-order programs, the obtained output programs resemble static-single-assignment (SSA) code. We emphasize the modularity of our approach and show how our language can easily be enriched with more optimized primitives, as required for some speed-ups in practice.
Explore related subjects
Keep this discovery
Mathieu Huot, Amir Shaikhha. 2022-12-19. Denotationally Correct, Purely Functional, Efficient Reverse-mode Automatic Differentiation. https://arxiv.org/abs/2212.09801
Cite the original work for its findings. Save a collection to share your selection of sources.