Reverse-mode differentiation is used for optimization, but it introduces references, which break the purity of the underlying programs, making them notoriously harder to optimize. We present a reverse-mode differentiation on a purely functional language with array operations. It is the first one to deliver a provably efficient, purely functional, and denotationally correct reverse-mode differentiation. We show that our transformation is semantically correct and verifies the cheap gradient principle. Inspired by PROPs and compilation to categories, we introduce a novel intermediate representation that we call 'unary form'. Our reverse-mode transformation is factored as a compilation scheme through this intermediate representation. We obtain provably efficient gradients by performing general partial evaluation optimizations after our reverse-mode transformation, as opposed to manually derived ones. For simple first-order programs, the obtained output programs resemble static-single-assignment (SSA) code. We emphasize the modularity of our approach and show how our language can easily be enriched with more optimized primitives, as required for some speed-ups in practice.
翻译:反向模式微分常用于优化,但其引入的引用会破坏底层程序的纯函数性,使得程序优化更加困难。本文提出了一种基于纯函数式语言(包含数组操作)的反向模式微分方法。这是首个能够提供可证明高效、纯函数式且形式语义正确的反向模式微分方法。我们证明了该变换在语义上是正确的,并验证了廉价梯度原理。受PROP(乘积与置换范畴)和范畴编译的启发,我们引入了一种新颖的中间表示——"一元形式"。通过该中间表示,我们将反向模式变换分解为一个编译方案。通过对变换后的程序执行通用部分求值优化(而非人工推导的优化),我们获得了可证明高效的梯度。对于简单的一阶程序,所得到的输出程序类似于静态单赋值(SSA)代码。我们强调该方法的模块化特性,并展示了如何通过添加更优化的原语来轻松扩展我们的语言——这在实践中对加速计算至关重要。