[论文解读] Zygote: A Differentiable Programming System to Bridge Machine Learning and Scientific Computing

Mike Innes, Alan Edelman+5 more|arXiv (Cornell University)|Jul 17, 2019

Parallel Computing and Optimization Techniques参考文献 39被引用 63

一句话总结

Zygote 是在 Julia 中实现的可微编程系统，能够对通用程序结构（包括控制流、递归和可变性）进行自动微分，且无需用户干预。它支持高性能的分阶段计算，并与科学计算库无缝集成，通过表达性强的可微编程，实现机器学习与科学计算的统一框架。

ABSTRACT

Scientific computing is increasingly incorporating the advancements in machine learning and the ability to work with large amounts of data. At the same time, machine learning models are becoming increasingly sophisticated and exhibit many features often seen in scientific computing, stressing the capabilities of machine learning frameworks. Just as the disciplines of scientific computing and machine learning have shared common underlying infrastructure in the form of numerical linear algebra, we now have the opportunity to further share new computational infrastructure, and thus ideas, in the form of Differentiable Programming. We describe Zygote, a Differentiable Programming system that is able to take gradients of general program structures. We implement this system in the Julia programming language. Our system supports almost all language constructs (control flow, recursion, mutation, etc.) and compiles high-performance code without requiring any user intervention or refactoring to stage computations. This enables an expressive programming model for deep learning, but more importantly, it enables us to incorporate a large ecosystem of libraries in our models in a straightforward way. We discuss our approach to automatic differentiation, including its support for advanced techniques such as mixed-mode, complex and checkpointed differentiation, and present several examples of differentiating programs.

研究动机与目标

• 通过在通用程序结构上支持可微编程，弥合机器学习与科学计算之间的鸿沟。
• 在 Julia 中以最小的用户开销支持自动微分，避免代码重构或分阶段计算的需求。
• 通过统一且表达性强的编程模型，实现科学计算库与机器学习模型的无缝集成。
• 将自动微分扩展至混合模式、复数和检查点微分等高级技术。
• 证明可微编程既具备高性能又具备通用性，能够支持现实世界中的科学与机器学习工作负载。

提出的方法

• Zygote 通过利用 Julia 的多分派和元编程能力，生成高效的梯度计算，从而实现高性能代码的编译。
• 它使用源到源转换技术，自动微分任意 Julia 程序，包括具有控制流、递归和可变状态的程序。
• 该系统采用反向模式自动微分引擎，通过对抽象语法树的符号操作，支持混合模式和复数微分。
• 通过选择性地重新计算中间值，实现检查点机制，以减少反向传播期间的内存使用。
• Zygote 与 Julia 现有的数值和科学计算库生态系统集成，实现复杂科学工作负载的端到端可微性。
• 该框架设计为可扩展，允许用户以极少样板代码定义自定义可微操作。

实验结果

研究问题

• RQ1 可微编程系统是否能够在不需代码重构的情况下，支持控制流、递归和可变性等通用编程构造？
• RQ2 可微编程系统在复杂科学计算工作负载上，其梯度的编译与执行效率如何？
• RQ3 在通用编程语言中，自动微分在多大程度上可以扩展以支持混合模式和复数微分？
• RQ4 可微编程是否能够无缝地将科学计算库集成到机器学习模型中？
• RQ5 在高级语言（如 Julia）中启用完整程序可微性时，会带来哪些性能与可用性权衡？

主要发现

• Zygote 能够在无需任何用户干预或代码重构的情况下，成功微分包含复杂控制流和可变性的通用 Julia 程序。
• 该系统通过为梯度计算生成高效、低级别的代码，实现高性能，大多数情况下避免了运行时开销。
• 通过程序抽象语法树的符号操作，原生支持混合模式和复数微分。
• 检查点机制被有效集成，在保持计算效率的同时减少了反向传播期间的内存消耗。
• 该框架支持直接在可微模型中使用科学计算库（如常微分方程求解器和线性代数库），显著提升了表达能力与可重用性。
• Zygote 证明了可微编程能够在一个高性能、用户友好的统一系统中，实现机器学习与科学计算的融合。