自从我们最初提出需要一个用于机器学习 (ML) 的一流语言、编译器和生态系统以来，该领域已经出现了许多有趣的发展。现有的系统（例如 TensorFlow 和 PyTorch）中的权衡不仅没有得到解决，而且现在这两个框架都包含不同的“静态图”和“急切执行”接口，使得这些权衡更加清晰。同时，ML 模型从根本上来说是可微算法——通常称为可微编程——这一理念也开始流行起来。
在当前框架存在不足的地方，一些激动人心的新项目已经涌现出来，它们完全摒弃了图的概念，将可微编程带入了主流。Theano 团队的Myia将 Python 的一个子集微分并编译成高性能 GPU 代码。Swift for TensorFlow扩展了 Swift，以便兼容函数可以编译成 TensorFlow 图。最后，Flux 生态系统正在使用许多以 ML 为中心的工具扩展 Julia 的编译器，包括一流的梯度、即时 CUDA 内核编译、自动批处理以及对 TPU 等新型硬件的支持。
所有这些项目都具有巨大的潜力，但我们认为 Julia 具有优势。这篇文章基于我们将在 NeurIPS MLSys 上发表的论文，将探讨我们如何利用 Julia 从头开始重新思考 ML 工具，并提供对现代 ML 工具需要完成的工作的一些见解。
我们需要一种语言来编写可微算法，而 Flux 将 Julia 作为这种语言。Julia 从一开始就设计用于数学和数值计算，因此非常适合表达 ML 算法。同时，它将现代设计与编译器中的新思想相结合，使其更容易满足尖端 ML 的高性能需求。
在典型的框架都是由数十万行 C++ 代码组成的包罗万象的整体的情况下，Flux 仅仅是一千行简单的 Julia 代码。只需取一个用于梯度的包 (Zygote.jl)、一个用于 GPU 支持的包 (CuArrays.jl)，撒上一些简单的便捷函数，烘烤 15 分钟，就会得到一个功能齐全的 ML 堆栈。
与其他下一代 ML 系统一样，Flux 致力于提供一个直观的（“急切”或“定义时运行”）接口，并且坚决反对任何形式的图构建或性能注释。我们支持该语言的所有功能，从控制流和数据结构到宏。用户可以在 Jupyter Notebook 中交互式地编写代码，并将高性能数值与方便的绘图和可视化相结合。但我们也希望获得传统“静态图”框架带来的好处——零开销源到源 AD、算子融合、多 GPU/分布式训练和单二进制部署。
我们如何做到这一切？实际上，我们需要直接从编写的 Julia 语法中提取和分析“静态图”，这实际上是编译器的正常工作。大多数 ML 系统问题最终都是标准的且经过充分研究的编译器问题，只是从正确的角度来看。使用编译语言足以解决许多问题，而扩展该编译器是解决更多问题的最佳方法。我们在这篇文章中仅介绍了我们当前在该领域工作的一个示例——即获取梯度、为 GPU 和 TPU 编译以及自动批处理。
在推动反向模式微分的极限时，我们开始将其视为一个语言级问题。微分是一种符号变换，这是编译器的领域。现有的框架通过跟踪（实际上是一种部分求值或抽象解释）来实现这一点。引入了一种新的张量类型，它记录所有执行的基本数学运算，从而生成一个图（或符号表达式），其中删除了宿主语言的控制流和数据结构。但是，这带来了一个难以权衡的问题：我们要么接受解释器的开销（急切执行），要么冻结用户控制流并限制可以构建的模型类型（静态图）。
如果“图”只是 Julia 本身的语法呢？将这个想法发挥到极致，我们构建了Zygote，它直接作用于 SSA 形式的 IR 并支持诸如控制流、递归、数据结构和宏之类的语言特性。然后，我们可以将生成的 SSA 形式的伴随代码通过LLVM等编译器，并获得传统编译器优化应用于我们的前向和后向传递的所有好处。此外，这种方法为使用更高级和特定于领域的优化（例如内核融合和编译到 TPU 等加速器）扩展编译器基础设施提供了机会。Swift for TensorFlow 和Myia的开发人员在源到源 AD 技术的复兴中也正在探索类似的方法。
Julia 用于此任务的一个关键优势是它可以用来实现基本的数值库，例如微分方程求解器或优化库；这巧妙地解决了 ML 社区中日益增长的需求，在该社区中，研究人员对高性能代码（例如光线追踪器和物理引擎）进行反向传播，但梯度仍然必须用 C++ 手动实现。相比之下，由于 Julia 的实现是用 Julia 编写的，因此从ODE到金融定价模型的一切都可以轻松地进行微分。将这些强大的工具引入模型是深度学习真正成为可微编程的地方。
GPU 编程是现代 ML 的重要组成部分。但 GPU 通常被视为一个实现细节；框架在内部提供内核，但用户只能看到有限的数学运算集，无法直接对 GPU 进行编程。相比之下，Julia 中的 GPU 编程从头到尾都是一流的，一直到 CUDA 内核（可以很方便地从脚本或 Notebook 中编写和运行）。
一个简单的向量加法内核看起来类似于 CUDA C 等效项。
function kernel_vadd(a, b, c) i = (blockIdx().x-1) * blockDim().x + threadIdx().x c[i] = a[i] + b[i] return end
但是，Julia 的类型专门化使 GPU 上能够使用一组强大的附加抽象。例如，上面的代码不限于浮点数的密集数组，而是可以给出复数的稀疏数组；Julia 的正常专门化机制会即时生成一组新的 PTX 指令。我们甚至可以将此代码进一步抽象成一个“高阶内核”，它接受+函数（或*，或任意用户定义的f），从而在四行代码中创建整个map(f, x, y)函数族。
这使得一些强大的技巧成为可能，即使您从未自己编写过 CUDA 代码。例如，我们可以透明地融合一个大的广播表达式，例如1 / (1 + exp(-x))，以及它的后向传递，到一个 GPU 内核中，从而获得显著的加速。我们预计原生 GPU 代码生成能力和生态系统将在未来为各种基于 Julia 的机器学习库提供动力。
更进一步，Google 最近公开了其 Cloud TPU 使用的 XLA IR，使其他框架和 ML 以外的用户能够利用这种重量级硬件。XLA 功能强大但有限：它无法运行 Python 解释器，当然也无法获得良好的性能。然后，框架最终会与梯度处于类似的位置——它们别无选择，只能使用程序跟踪来剔除 Python，并最终得到一种快速但功能更有限的 ML 语言。
我们的回应是可以预见的：我们只需要从编写的 Julia 程序中提取“静态图”并将其直接编译到 XLA，从而允许 Julia 本身在 TPU 上运行。（实际上，这只是 Julia 通常编译过程的一个简单扩展，该过程会在将程序发送到 LLVM 之前从程序中提取尽可能大的“静态子图”。）这使我们能够充分利用 Julia 语言的表达能力，包括控制流、递归、多重分派、高阶函数、强大的数据结构和抽象、自定义数值类型以及现有的包，例如微分方程求解器和线性代数例程。所有这些都在利用 TPU 内的高性能 systolic 阵列引擎带来的好处的同时运行。您可以立即尝试，其中包含ResNet 等大型 ML 模型和TSVD 等线性代数例程的示例。
为了最大限度地利用这些加速器——它们每次内核启动的开销可能很大，但随着输入大小的增加而扩展得非常好——通常需要批处理程序，同时对多个训练示例应用前向和后向传递。在简单的情况下，例如卷积网络，可以通过沿着额外的批处理维度连接 10 张图像来轻松处理此问题。但是，在处理结构可变的输入（例如树或图）时，此任务变得更加困难。
大多数研究人员通过手动承担批处理代码的重大负担来解决这个问题。已经为不同的框架提出了不同的解决方案（DyNet、TensorFlow Fold），这些解决方案在可能的情况下试探性地将一些高级操作批处理在一起，但这些解决方案通常要么存在自身可用性问题，要么无法达到手动编写的代码的性能。
我们认为这个问题与单程序多数据（SPMD）编程相同，后者在语言和编译器社区中已经被广泛研究了数十年，并在最近的批量处理方法（如matchbox）中变得更加明显。事实上，它与GPU内部使用的并行模型非常相似，并且已被实现为CPU SIMD单元的编译器转换。受此工作的启发，我们正在Julia中实现相同的转换，以提供用于标量SIMD单元和模型级批处理的SPMD编程。这使我们能够实现编写操作单个样本的简单代码的理想目标，同时仍然获得现代硬件的最佳性能。
我们相信机器学习的未来在于语言和编译器技术，特别是扩展新语言或现有语言以满足机器学习研究的高需求。这对机器学习社区和数值编程领域都有好处；能够很好地支持微分、矢量化和异构硬件的语言将足够强大，足以推动科学的许多进步。
在这些下一代工具——Myia、Swift/TF和Flux——达到与现有的框架（TensorFlow、PyTorch和Knet）一样生产就绪的程度之前，还有一段路要走。但是，如果您正在机器学习领域开拓新方向，它们很可能是您的最佳选择。尝试一下，看看机器学习的未来是什么样子。