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

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

Neural Networks and Applications被引用 129

一句话总结

本文概述 ∂P，是一个用 Julia 构建的可微分编程系统，使得可以对通用 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.

研究动机与目标

• 激发一个跨越 ML 与科学计算的统一可微分编程方法。
• 提供一个基于 Julia 的系统，能够以最少的用户干预对通用程序进行求导。
• 展示将机器学习与现有科学包相结合的实际应用（例如光线追踪、量子仿真、金融）。
• 展示具有自定义数据类型和复杂控制流的性能特征与可扩展性。

提出的方法

• 使用源到源自动微分方法（Zygote）从 Julia 代码生成梯度。
• 定义一个通用微分算子 J，它用雅可比-向量积将函数输出封装，以用于反向传播。
• 提供一个用户可扩展的偏导接口，通过多重分派实现自定义梯度。
• 在混合语言和复杂数据类型上演示可微分求导，而无需对包的原语进行更改。

实验结果

研究问题

• RQ1 可微分编程系统是否能够对具有控制流、变异和复杂数据类型的通用 Julia 程序进行求导？
• RQ2 可微分编程方法在 Julia 中与现有的 ML 和科学计算生态系统的集成程度如何？
• RQ3 在 CPU、GPU、TPU 上对大型真实世界 Julia 程序进行求导的性能影响是什么？
• RQ4 在保持生态系统可组合性的同时，可微分编程是否能够实现诸如可微分光线追踪、量子机器学习和神经 SDEs 等应用？

主要发现

• Zygote 通过源到源转换实现低开销自动微分，并在前向与反向传播之间重用信息。
• 对于堆叠的 LSTM，按对偶操作的开销平均约为 568.8 纳秒，与 PyTorch 的开销具有竞争力。
• 在各类领域（光线追踪、量子仿真、神经 SDE、金融）中通过 ∂P 在 Julia 中进行求导是可行且有成效的。
• 可以通过偏导接口为新类型轻松定义自定义梯度，而无需修改底层的 AD 系统。
• 该方法支持在 CPU、GPU、TPU 上的训练，在深度学习任务中与现有 ML 框架的性能相匹配。