[论文解读] Don't Unroll Adjoint: Differentiating SSA-Form Programs

Michael Innes|arXiv (Cornell University)|Oct 18, 2018

Parallel Computing and Optimization Techniques参考文献 22被引用 61

一句话总结

本文提出了一种基于 SSA 形式程序的反向模式自动微分方法，在 Julia 中实现为 Zygote，能够高效地对复杂语言特性进行微分并与编译器交互。

ABSTRACT

This paper presents reverse-mode algorithmic differentiation (AD) based on source code transformation, in particular of the Static Single Assignment (SSA) form used by modern compilers. The approach can support control flow, nesting, mutation, recursion, data structures, higher-order functions, and other language constructs, and the output is given to an existing compiler to produce highly efficient differentiated code. Our implementation is a new AD tool for the Julia language, called Zygote, which presents high-level dynamic semantics while transparently compiling adjoint code under the hood. We discuss the benefits of this approach to both the usability and performance of AD tools.

研究动机与目标

• 激发对超越基于追踪的方法、具表达性且高性能的自动微分的需求。
• 提出基于 SSA 形式的自动微分，能够与现有编译器集成，生成优化的导数代码。
• 展示 Zygote 的实现及其与 Julia 编译器和运行时的交互。
• 展示该方法如何支持如控制流、变更和高阶导数等语言特性。

提出的方法

• 通过 J 运算符和反向传播将梯度传播到 SSA 形式中间表示（IR）来表示微分。
• 在伴随中使用反向基本块和 phi 节点来处理控制流和数据依赖。
• 在栈上存储并解析 alpha 值，以在反向微分过程中管理 SSA 迭代定义。
• 为列表单元、箱子和闭包等复杂数据结构及变更提供相应的伴随反向传播（pullbacks）。
• 利用 Julia 的类型推断和编译器优化来内联 pullbacks 并生成高效代码。
• 将该方法定位为高层次可微编程与传统编译器优化之间的桥梁。

实验结果

研究问题

• RQ1 基于 SSA 形式 IR 的反向模式 AD 是否能够容纳包括控制流、变更和高阶函数在内的完整语言特性？
• RQ2 对 SSA 形式程序进行微分如何与宿主编译器交互以生成高效的微分代码？
• RQ3 与基于追踪的 AD 系统相比，性能和可用性之间的权衡是什么？
• RQ4 如何生成伴随代码以在分支和循环中保持正确的数据流？
• RQ5 该方法是否可扩展到现实世界的模型和编译后端（如 LLVM、TPUs）？

主要发现

• Zygote 提供了一个可工作的用于 Julia 的 SSA 形式 AD 工具，能够产生高效的微分代码。
• 伴随构造通过使用 phi 节点和 alpha 跟踪，在反向 CFG 中保持正确的数据流。
• 反向传播（pullbacks）可以组合以支持高阶导数和包括变更以及数据结构在内的复杂语言特性。
• 实验在若干场景下显示出的性能与手写导数竞争，并且对比追踪方法的基准结果较有利。
• 该方法使得能够利用现有编译器优化和未来的加速器，同时不牺牲表达性。
• 基于 SSA 的微分可以与现代编译器栈整合，实现跨语言的微分编程。