Provably correct, asymptotically efficient, higher-order reverse-mode automatic differentiation

Krawiec, Faustyna; Jones, Simon Peyton; Krishnaswami, Neel; Ellis, T. M. R.; Eisenberg, Richard A.; Fitzgibbon, Andrew

doi:10.1145/3498710

Cited by 21 publications

(15 citation statements)

References 12 publications

Supporting

Mentioning

Contrasting

Order By: Relevance

“…Furthermore, reverse-mode automatic differentiation of control flow constructs (including higher-order functions) cannot be performed without forming a dynamic trace of the computation, and the languages conventionally used in those traces have expressive power similar to Linear A. For example Pearlmutter and Siskind [2008a]; Wang et al [2019] represent the trace as a sequence of applications, while Krawiec et al [2022] use a language almost identical to the linear fragment of Linear A.…”

Section: Languagementioning

confidence: 99%

“…Mazza and Pagani [2021] outline semantics for a higher-order language with conditionals that are consistent with the mathematical notion of differentiation everywhere except for a measure 0 set. Finally, Krawiec et al [2022] outline a provably correct and efficient implementation of AD in a higher-order language, at a cost of requiring a monadic translation.…”

Section: Related Workmentioning

confidence: 99%

See 1 more Smart Citation

You Only Linearize Once: Tangents Transpose to Gradients

Radul¹,

Paszke²,

Frostig³

et al. 2022

Preprint

View full text Add to dashboard Cite

Automatic differentiation (AD) is conventionally understood as a family of distinct algorithms, rooted in two "modes"-forward and reverse-which are typically presented (and implemented) separately. Can there be only one? Following up on the AD systems developed in the JAX and Dex projects, we formalize a decomposition of reverse-mode AD into (i) forward-mode AD followed by (ii) unzipping the linear and non-linear parts and then (iii) transposition of the linear part.To that end, we use the technology of linear types to formalize a notion of structurally linear functions, which are then also algebraically linear. Our main results are that forward-mode AD produces structurally linear functions, and that we can unzip and transpose any structurally linear function, conserving cost, size, and structural linearity. Composing these three transformations recovers reverse-mode AD. This decomposition also sheds light on checkpointing, which emerges naturally from a free choice in unzipping let expressions. As a corollary, checkpointing techniques are applicable to general-purpose partial evaluation, not just AD.We hope that our formalization will lead to a deeper understanding of automatic differentiation and that it will simplify implementations, by separating the concerns of differentiation proper from the concerns of gaining efficiency (namely, separating the derivative computation from the act of running it backward).

show abstract

Section: Languagementioning

confidence: 99%

Section: Related Workmentioning

confidence: 99%

You Only Linearize Once: Tangents Transpose to Gradients

Radul¹,

Paszke²,

Frostig³

et al. 2022

Preprint

View full text Add to dashboard Cite

show abstract

“…There has recently been a flurry of work studying AD from a programming language point of view, a lot of it focussing on functional formulations of AD and their correctness. Examples of such papers are [24,9,27,6,1,13,21,34,19,14,35,16,29]. Of these papers, [24,1,21,29] are particularly relevant as they also consider automatic differentiation of languages with partial features.…”

Section: Related Workmentioning

confidence: 99%

“…Dual numbers techniques give very simple forward and reverse mode AD algorithms for expressive ML-family functional languages [27,6,13,21,16,29]. The correctness proofs for these algorithms rely on (open) semantic logical relations arguments.…”

Section: Introductionmentioning

confidence: 99%

Logical Relations for Partial Features and Automatic Differentiation Correctness

Nunes¹,

Vákár²

2022

Preprint

View full text Add to dashboard Cite

We present a simple technique for semantic, open logical relations arguments about languages with recursive types, which, as we show, follows from a principled foundation in categorical semantics. We demonstrate how it can be used to give a very straightforward proof of correctness of practical forward-and reverse-mode dual numbers style automatic differentiation (AD) on ML-family languages. The key idea is to combine it with a suitable open logical relations technique for reasoning about differentiable partial functions (a suitable lifting of the partiality monad to logical relations), which we introduce.

show abstract

“…Notably, (Brunel et al, 2020) shows that their algorithm has the right complexity if one assumes a specific operational semantics for their linear -calculus with what they call a "linear factoring rule". Very recently, (Krawiec et al, 2021) applied the idea of reverse AD through tracing to a higher-order functional language with variant types. They implement the custom operational semantics as an evaluator and give a denotational correctness proof (using logical relations techniques similar to those of (Barthe et al, 2020;Huot et al, 2020)) as well as an asymptotic complexity proof about the full code transformation plus evaluator.…”

Section: Related Workmentioning

confidence: 99%

CHAD for Expressive Total Languages

Nunes¹,

Vákár²

2021

Preprint

View full text Add to dashboard Cite

We show how to apply forward and reverse mode Combinatory Homomorphic Automatic Differentiation (CHAD) (Vákár, 2021 ,b) to total functional programming languages with expressive type systems featuring the combination of • tuple types; • sum types; • inductive types; • coinductive types; • function types.We achieve this by analysing the categorical semantics of such types in Σ-types (Grothendieck constructions) of suitable categories. Using a novel categorical logical relations technique for such expressive type systems, we give a correctness proof of CHAD in this setting by showing that it computes the usual mathematical derivative of the function that the original program implements. The result is a principled, purely functional and provably correct method for performing forward and reverse mode automatic differentiation (AD) on total functional programming languages with expressive type systems.

show abstract

Provably correct, asymptotically efficient, higher-order reverse-mode automatic differentiation

Cited by 21 publications

References 12 publications

You Only Linearize Once: Tangents Transpose to Gradients

You Only Linearize Once: Tangents Transpose to Gradients

Logical Relations for Partial Features and Automatic Differentiation Correctness

CHAD for Expressive Total Languages

Contact Info

Product

Resources

About