arXiv · 2105.09469
Decomposing reverse-mode automatic differentiation
Abstract
We decompose reverse-mode automatic differentiation into (forward-mode) linearization followed by transposition. Doing so isolates the essential difference between forward- and reverse-mode AD, and simplifies their joint implementation. In particular, once forward-mode AD rules are defined for every primitive operation in a source language, only linear primitives require an additional transposition rule in order to arrive at a complete reverse-mode AD implementation. This is how reverse-mode AD is written in JAX and Dex.
Explore related subjects
Keep this discovery
Roy Frostig, Matthew J. Johnson, Dougal Maclaurin, Adam Paszke, Alexey Radul. 2021-05-20. Decomposing reverse-mode automatic differentiation. https://arxiv.org/abs/2105.09469
Cite the original work for its findings. Save a collection to share your selection of sources.