一句话:将 JAX Autodiff 编码到线性 λ-演算,证明定性与定量可靠性,并说明 unzipping 可选。
JAX Autodiff 是 JAX、Dex 等系统中自动微分的核心,此前 Radul 等人(POPL 2023)已用线性类型演算形式化,但该演算高度专用,其类型系统是否对应一个有独立意义的子结构逻辑尚不明确。本文把 JAX Autodiff 编码进与 Girard 线性逻辑具有 Curry-Howard 对应的线性 λ-演算,并证明编码在定性上保持外延等价,在定量上保持 Radul 等人定义的工作代价。作为副产品,作者表明在 JAX Autodiff 中用于实现反向传播的变换 unzipping 实际上是可选的。
JAX Autodiff refers to the core of the automatic differentiation (AD) systems developed in projects like JAX and Dex. JAX Autodiff has recently been formalised in a linear typed calculus by Radul et al in POPL 2023. Although this formalisation suffices to express the main program transformations of AD, 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 JAX 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 by Radul et al. As a byproduct, we show that unzipping, one of the transformations used to implement backpropagation in JAX Autodiff, is, in fact, optional.