Year: 2021
Type: preprint
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.
Source: arXiv (Cornell University)
Authors: Frostig, Roy, Johnson, Matthew J., Maclaurin, Dougal, Paszke, Adam, Radul, Alexey
Language: English
Subfield: Artificial Intelligence
Field: Computer Science
Domain: Physical Sciences
Open Access status: green
OpenAlex ID: W3162555546