一次线性化足矣：正切转置为梯度
编程语言
2022-12-07 v2
摘要
自动微分（AD）传统上被理解为一系列不同的算法，根植于两种“模式”——前向模式与反向模式——它们通常被分开呈现（并实现）。能否只有一种？继 JAX 与 Dex 项目中开发的 AD 系统之后，我们形式化地将反向模式 AD 分解为 (i) 前向模式 AD，随后 (ii) 将线性与非线性部分解压（unzip）并 (iii) 对线性部分进行转置。为此，我们定义了一种（子结构上的）线性类型系统，可证明一类函数是（代数上）线性的。我们的主要结果是：前向模式 AD 产生此类线性函数，且我们可对任何此类线性函数进行解压与转置，同时保持代价、规模与线性性。组合这三种变换即可恢复反向模式 AD。该分解也阐明了检查点技术（checkpointing），它自然地从解压 `let` 表达式时的自由选择中浮现。作为一个推论，检查点技术可应用于通用部分求值，而不仅限于 AD。我们希望我们的形式化能促进对自动微分更深层次的理解，并通过将微分本身的问题与获取效率的问题（即，将导数计算与反向运行它的动作分离）分开，从而简化实现。
引用
@article{arxiv.2204.10923,
title = {You Only Linearize Once: Tangents Transpose to Gradients},
author = {Alexey Radul and Adam Paszke and Roy Frostig and Matthew Johnson and Dougal Maclaurin},
journal= {arXiv preprint arXiv:2204.10923},
year = {2022}
}