JAX Autodiff from a Linear Logic Perspective (Extended Version)

TL;DR

This work formalizes JAX-style automatic differentiation by encoding Linear A, the domain-specific calculus for Autodiff, into a linear λ-calculus (λLL) with a Curry–Howard correspondence to Girard's linear logic. The authors define δ, a translation from Linear A to λLL, and three program transformations—Forward (F), Unzipping (U), and Transpose (T)—that replicate forward-mode evaluation, tape-based tangent propagation, and the adjoint gradient computation, respectively. They prove qualitative and quantitative soundness: δ faithfully represents Linear A terms up to extensional equivalence, and F, U, T preserve the original work cost with the option to skip unzipping without changing results. The framework highlights the modularity of backpropagation, enabling flexible compositions of forward and transpose passes and offering a rigorous baseline to compare AD formalisms with backpropagators (e.g., POPL2020). The work’s significance lies in grounding AD in a substructural logic, yielding a principled, workload-aware account of backpropagation that could influence the design of AD in libraries like JAX and Dex.

Abstract

Autodiff refers to the core of the automatic differentiation systems developed in projects like JAX and Dex. Autodiff has recently been formalised in a linear typed calculus by Radul et al in arXiv:2204.10923. Although this formalisation suffices to express the main program transformations of Autodiff, the calculus is very specific to this task, and it is not clear whether the type system yields a substructural logic that has interest on its own. We propose an encoding of Autodiff into a linear $λ$-calculus that enjoys a Curry-Howard correspondence with Girard's linear logic. We prove that the encoding is sound both qualitatively (the encoded terms are extensionally equivalent to the original ones) and quantitatively (the encoding preserves the original work cost as described in arXiv:2204.10923). As a byproduct, we show that unzipping, one of the transformations used to implement backpropagation in Autodiff, is, in fact, optional.

Figures (30)

• Figure 1: Linear A Typing Rules. For short, we detail only the case of $\underline{f}$ binary, the general case being immediate.
• Figure 2: Forward of some purely primal Linear B expressions. In the case of a $\mathtt{let}$ expression we decompose the set its free variables into $dom(\phi_i) = FV^{}\!(e^{p}_i) \setminus (FV^{}\!(e^{p}_1) \cap FV^{}\!(e^{p}_2))$ and $\{z_i\}^k_{i=1} = FV^{}\!(e^{p}_1) \cap FV^{}\!(e^{p}_2)$.
• Figure 3: Some cases of unzipping. In the case of a let expression, we suppose $\mathcal{U}^{\mathtt{Jax}}(e_i)= E_i \texttt{ in } (e_i^{p};\dot{e_i})$.
• Figure 4: Transpose of some purely tangent Linear B expressions.
• Figure 5: Application of JAX Autodiff to the Linear B expression computing the map $(sin(x)*y)+cos(x)$. Color blue (resp. red) underlines the tangent (resp. tangent transpose) subexpressions introduced by JAX Autodiff.

Theorems & Definitions (98)

• Remark 1
• Remark 2
• Remark 2
• theorem 1: Subject Reduction
• theorem 2: Strong Normalisaton
• theorem 3: Confluence
• Definition 1: $\sim_A$, $\sim_{\Gamma\vdash A}$
• Definition 2: Safe term
• Proposition 1
• Remark 3