[论文解读] Fashionable Modelling with Flux

Michael Innes, Elliot Saba+7 more|arXiv (Cornell University)|Nov 1, 2018

Computer Graphics and Visualization Techniques参考文献 21被引用 135

一句话总结

Flux 是一个基于 Julia 的可微分编程框架，强调简单性、可 hack 与编译器驱动加速，能够实现灵活的模型并便于与加速器的集成。

ABSTRACT

Machine learning as a discipline has seen an incredible surge of interest in recent years due in large part to a perfect storm of new theory, superior tooling, renewed interest in its capabilities. We present in this paper a framework named Flux that shows how further refinement of the core ideas of machine learning, built upon the foundation of the Julia programming language, can yield an environment that is simple, easily modifiable, and performant. We detail the fundamental principles of Flux as a framework for differentiable programming, give examples of models that are implemented within Flux to display many of the language and framework-level features that contribute to its ease of use and high productivity, display internal compiler techniques used to enable the acceleration and performance that lies at the heart of Flux, and finally give an overview of the larger ecosystem that Flux fits inside of.

研究动机与目标

• 激励一个可编程的 ML 框架，将模型视为用高级语言表达的可微分算法。
• 演示 Flux 的设计选择，使在 Julia 中-building 简单、可扩展的模型成为可能。
• 展示面向 GPU、TPU 和部署的编译器驱动性能特性。
• 强调 Flux 如何与更广泛的 Julia 生态系统整合，以实现新颖的 ML 工作流。

提出的方法

• 将 Flux 呈现为在 Julia 中构建的精简、高层 ML 堆栈。
• 解释三大支柱：简单性、可 hack 与编译器技术。
• 描述 Zygote 自源到自源的自动微分器及其基于 LLVM 的优化管线。
• 展示 Flux 如何通过编译器基础设施实现 GPU/TPU 编译和自动批处理。
• 提供诸如 Discriminative Adversarial Networks 等示例，以说明灵活的训练范式。
• 讨论与其他 Julia 包在优化、微分方程和概率编程方面的互操作性。

实验结果

研究问题

• RQ1 相较于传统框架，Flux 如何在实现 ML 模型方面实现简单性？
• RQ2 Flux 如何通过允许自定义层、梯度和训练循环而不牺牲性能来实现 hackability？
• RQ3 使用了哪些编译器技术来在 GPU/TPU 上加速 Flux 模型，并在动态语言中实现静态化的优化？
• RQ4 Flux 如何在可微分编程工作流中利用 Julia 生态系统（AD、微分方程、优化）？
• RQ5 具体的模型范式（例如 Discriminative Adversarial Networks）有哪些，能展示 Flux 的灵活性？

主要发现

• Flux 为核心功能实现了上千行 Julia 代码，从而实现高开发者生产力。
• Flux 支持在单个模型中混合内置和用户定义的层、算法以及控制流。
• Zygote 基于编译器的 AD 方法实现了语言级微分，并带有 LLVM 优化。
• Flux 能将 Julia 代码编译到 GPU/TPU，包括将整个训练循环编译到 XLA 以用于 TPU。
• 通过受编译器启发的 SPMD 转换实现自动批处理，适用于变结构输入的可扩展批处理。
• Flux 提供一个精简的集成层，与 GPU 数组及其他 Julia 包协同工作，构建一个有机的 ML 生态系统。