Decomposing reverse-mode automatic differentiation
Programming Languages
2021-05-21 v1 Machine Learning
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.
Cite
@article{arxiv.2105.09469,
title = {Decomposing reverse-mode automatic differentiation},
author = {Roy Frostig and Matthew J. Johnson and Dougal Maclaurin and Adam Paszke and Alexey Radul},
journal= {arXiv preprint arXiv:2105.09469},
year = {2021}
}
Comments
Presented at the LAFI 2021 workshop at POPL, 17 January 2021