JAX Autodiff from a Linear Logic Perspective (Extended Version)

2025-10-19

AI Summary

Automatic differentiation (Autodiff) lacks a unifying logical foundation. Method: We encode Autodiff into a linear λ-calculus endowed with Curry–Howard correspondence, thereby embedding it within Girard’s linear logic framework; specifically, we formalize the core Autodiff primitives—forward and reverse mode propagation—as program transformations in a linear type system, and rigorously verify semantic equivalence and computational complexity preservation. Contributions: (1) We establish, for the first time, that Autodiff admits a faithful interpretation as a substructural logic with independent logical meaning; (2) we demonstrate that the “unzipping” operation commonly assumed essential in reverse-mode Autodiff is logically redundant, enabling simpler, more modular implementations; (3) our approach achieves significantly enhanced theoretical rigor and system concision—without incurring additional computational overhead—providing the first semantics of Autodiff grounded in linear logic.

Technology Category

• Knowledge Representation and Reasoning: Logic Programming
• Constraint Satisfaction and Optimization: Satisfiability Modulo Theories
• Reasoning under Uncertainty: Probabilistic Programming

Application Category

• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Systems and Infrastructure for Web, Mobile and WoT: Applied ML and AI for Web-based mobile applications

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.

Problem

• Encoding Autodiff into linear λ-calculus with Curry-Howard correspondence
• Proving soundness of encoding for extensional equivalence and cost
• Demonstrating optionality of unzipping transformation in backpropagation

Innovation

• Encodes Autodiff into linear lambda calculus
• Links calculus to Girard's linear logic
• Shows unzipping transformation is optional