一句话:用稀疏向量、状态传递和去函数化优化 CHAD,使反向模式自动微分达到预期复杂度,并给出 Agda 复杂度证明。
本文研究如何让基本组合同态自动微分(CHAD)在保持简单、可组合和通用性的同时,达到反向模式自动微分应有的计算复杂度。作者组合稀疏向量、状态传递风格代码以及高阶语言中的去函数化/闭包转换等已知优化,得到纯函数式算法,并用(函数式)可变更新消除最后的对数因子。作者在 Agda 中形式化了复杂度证明。论文还讨论将该技术用于并行函数式数组程序:所需可变性均为可交换、可结合的累加,从而保留任务并行性,并且可为大多数数据并行数组原语写出数据并行导数。
We show how the basic Combinatory Homomorphic Automatic Differentiation (CHAD) algorithm can be optimised, using well-known methods, to yield a simple, composable, and generally applicable reverse-mode automatic differentiation (AD) technique that has the correct computational complexity that we would expect of reverse-mode AD. Specifically, we show that the standard optimisations of sparse vectors and state-passing style code (as well as defunctionalisation/closure conversion, for higher-order languages) give us a purely functional algorithm that is most of the way to the correct complexity, with (functional) mutable updates taking care of the final log-factors. We provide an Agda formalisation of our complexity proof. Finally, we discuss how the techniques apply to differentiating parallel functional array programs: the key observations are 1) that all required mutability is (commutative, associative) accumulation, which lets us preserve task-parallelism and 2) that we can write down data-parallel derivatives for most data-parallel array primitives.