Decomposing reverse-mode automatic differentiation

20 May 2021  ·  Roy Frostig, Matthew J. Johnson, Dougal Maclaurin, Adam Paszke, Alexey Radul ·

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.

PDF Abstract

Code
Add Remove Mark official

No code implementations yet. Submit your code now

Tasks
Add Remove

Datasets

  Add Datasets introduced or used in this paper

Results from the Paper
Add Remove

  Submit results from this paper to get state-of-the-art GitHub badges and help the community compare results to other papers.

Methods
Add Remove

No methods listed for this paper. Add relevant methods here
Contact us on: hello@paperswithcode.com . Papers With Code is a free resource with all data licensed under CC-BY-SA.
Terms Data policy Cookies policy