梯度下降是最好的程序员：Julia未来将内嵌可微编程系统

如果一种语言，它本身就是一个可微编程系统，那写起代码来又是怎样的。

近日，Julia Computing 团队发表论文表示他们构建了一种可微编程系统，它能将自动微分内嵌于 Julia 语言，从而将其作为第一级的语言特性。也就是说，我们以后直接用 Julia 语言及可微编程就能写模型了？都不需要再调用 TensorFlow 或 PyTorch 这样的框架了？

如果我们将可微编程系统视为编程语言最重要的特性之一，那么不论是机器学习还是其它科学计算都将方便不少，这样的语言也将是科学计算最好的语言。Y Combinator Research 研究者 Michael Nielsen 对此也非常兴奋，他非常赞同 Andrej Karpathy 所说的「梯度下降是更好的程序员」。

Karpathy 同时也回复到：「我们正向前移动了一点，与原来对程序有完整的定义不同，我们现在只是写一个大致的架构，这样的架构会通过权重参数化。如果我们有一个好的评估标准，那么最优化算法就能帮我们找到更好的解。」

这里 Karpathy 说的就是机器学习和编程的区别，ML 会通过梯度下降等最优化方法自动搜索最优解。但这里有个问题，模型需要梯度才能向着最优前进，因此模型的很多部分都要求是可微的。鉴于这一点，很多人也就将 ML 称呼为可微编程了。

但是可微编程只能用于机器学习吗？它能不能扩展到其它领域，甚至成为编程语言的基本特性？答案是可以的，这就是 Julia 团队及 MIT 等其他研究机构正在尝试的。

近年来，机器学习模型越来越精妙，展现出了很多科学计算的特性，侧面凸显了机器学习框架的强大能力。研究者表示，由于广泛的科学计算和机器学习领域在底层结构上都需要线性代数的支持，因此有可能以可微编程的形式，创造一种新的计算基础设施。

• 论文地址：https://arxiv.org/pdf/1907.07587.pdf

• GitHub项目地址：https://github.com/MikeInnes/zygote-paper

可微编程与 DL 框架又有什么不同

在论文中，研究人员提出了一种可微编程系统，它能在 Julia 语言中完成梯度计算，并成为 Julia 语言的第一级特性。这使得构建深度学习模型变得更直观。更重要的是，这令用户可以使用已有的 Julia 科学计算包去构建深度学习模型，并高效实现梯度计算。

那么可微编程与 TensorFlow 或 PyTorch 等 DL 框架又有什么不同，它的速度难道还要快于这些构建细算图的系统？研究者表示，可微编程会执行 source-to-source 的转换，自动微分转换基本上没有运行时开销，因此它要比反向传播的实际计算成本更有优势。

在 Reddit 上也有开发者讨论到了这一点：

一位名为「Coconut_island」开发者表示，Zygote 和现有的机器学习框架的不同之处在于，Zygote 将源代码视为计算图，省略了中间语言的过程。更重要的是，通过控制流进行子梯度计算，这种梯度计算的方法可以应用在任何一种控制流分支上，具有代码的抽象性。

科学计算和机器学习的联姻： 可微编程

咋一看，科学计算和机器学习分属不同的领域。现代机器学习在神经网络上的突破使得整个学科取得了极大的进步。神经网络可以解决很多不同的计算机科学问题，使得人们开始设计新的硬件和软件，以便于提升性能，处理极大量的标注数据，并能够同时将训练好的模型部署在设备上。科学计算与之相反，它拥有非常悠久的历史，且希望使用一系列从物理现象中获取的建模技术。

和典型的机器学习研究者不同，许多科学计算学家在数据量较少、但计算复杂度更高和范围更广的数据上进行研究。但是观察来看，科学计算和机器学习有共同点。两个领域都倾向于使用动态编程语言进行计算，如 Python、R 和 Julia。通常来说，Python 和 R 的关键性能模块都使用 C++和 Fortran，而 Julia 中则较少。而且两者的核心计算流程都是基于线性代数的，并且有硬件专门用于加速这种计算。

由于机器学习和科学计算都依赖线性代数，且已有大量硬件专门优化用于加速运算，因此可微编程具有进一步融合两者的潜力。使用何种语言完成这项工作呢？由于编程工作量的问题，研究人员仅选择在 Julia 语言上增加可微编程能力。选择这门语言的一个原因是，Julia 语言已有大量的机器学习和科学计算包，都是纯粹的 Julia 语言实现。这可以帮助他们在相对较多的应用中进行测试。

研究人员表示，已有的 Julia 程序包可以直接使用他们的系统，包括处理用户定义类、基于状态的控制流，以及通过「source-to-source」自动微分进行大量标量运算。论文展示了一些使用可微编程的案例。

用 Zygote 对正弦函数进行微分求导

研究人员以正弦函数为例，解释可微编程的方法。如图为正弦函数 sin(X) 的泰勒展开：

这个公式可以使用 Julia 语言进行定义，为了运行，请安装 Julia 语言 1.1 以上版本，并安装 Zygote.jl 和 ForwardDiff.jl 包，代码如下所示：

using Pkg
Pkg.add("Zygote")
Pkg.add("ForwardDiff")
using Zygote, ForwardDiff
function s(x)
    t = 0.0
    sign = -1.0
    for i in 1:19
        if isodd(i)
            newterm = x^i/factorial(i)
            abs(newterm)<1e-8 && return t
            println("i=",i)
            sign = -sign
            t += sign * newterm
         end
         
    end
    return t
end

尽管用 Julia 语言可以更好的表示正弦函数，但是为了展示计算过程，这里使用了一个循环、一个条件语句、「isodd」和「factorial」函数，这些都是 Julia 的原生实现。在微分计算中，自动微分直接工作。如下为 x = 1.0 开始，i 的变化情况，计算得到的梯度，以及何时这个梯度和 cos(1.0) 的梯度匹配。

julia> ForwardDiff.derivative(s, 1.0) # Forward Mode AD
i=1
i=3
i=5
i=7
i=9
i=11
0.540302303791887

julia> Zygote.gradient(s, 1.0) # Reverse Mode AD
i=1
i=3
i=5
i=7
i=9
i=11
(0.5403023037918872,)

julia> cos(1.0)
0.5403023058681398

可微编程中的深度学习

Zygote 是用来计算深度学习模型梯度的灵活后端，它的效率非常高。下面展示了一个典型的例子，研究人员使用 LSTM 学习莎翁的作品。下面的代码展示了 Zygote 很多优秀的属性，只要调用几个 Julia 语言的一些便捷特性就能完成。

首先，定义的模型并没有特殊的数据类型，因此可以直接启用自动微分（AD）；模型的定义只需要确定前向传播的计算流就行了，反向传播只有在定义 BLAS 运算和基本数组操作等基本构建块时才需要手动定义。

Zygote 还可以用于包装模型损失的计算，明确指明计算边界应该对于模型梯度的计算是可微的，但是其它部分的代码，包括我们定义的 LSTM 层级，在写的时候都不需要考虑自动微分过程。最后模型可以在 CPU、GPU 和谷歌的 TPU 上运行，它们不需要或只需要一点修改。

# Load data and alphabet
alphabet, Xs, Ys = load_data("shakespeare_input.txt")

# Define simple LSTM-based model to map from alphabet back on to alphabet,
# predicting the next letter in a corpus of Shakespeare text.
model = Chain(
    LSTM(length(alphabet), 128),
    LSTM(128, 128),
    Dense(128, length(alphabet)),
    softmax,
)

opt = ADAM(0.01)

# Run through our entire dataset a few times
 for epoch_idx in 1:10,
    (x_batch, y_batch) in zip(Xs, Ys) 
    
    # Calculate gradients upon the model for this batch of data,
    # summing crossentropy loss across each time index in this batch
    grads = |\Zygoteplain|.gradient(model) do model
       return sum(crossentropy.(model.(x_batch), y_batch)) 
    end
    
    # Update the model, then reset its internal LSTM states
    model = update!(opt, model, grads)
    Flux.reset!(model)
end

如上所示整个模型及训练过程看起来非常简洁，有点像 Keras 或 PyTorch 的编程风格，基本上只需要搭建模型、算梯度、更新参数等几个重要步骤。即使用户不太了解 Julia 的语法规则，但还是很容易读懂的，它非常重要的部分就是通过 Zygote 求梯度。

Zygote 提供了一个成本极其低的自动微分接口。通过执行 source-to-source 的转换，AD 转换基本上没有运行时开销，因此它要比反向传播的实际计算成本更有优势。此外，Zygote 在 TPU pod 上训练 ResNet 与 TensorFlow 有相同的性能水平。

表 1： 在不同深度的 LSTM 网络中，Zygote 每一个运算（OP）开销的估计值。

为了度量上面的结果，研究者对 LSTM 网络的反向传播做了一个基准测试，并研究运行时间与批量大小的关系，从而进一步估计自动微分系统中每一个操作的固定开销。研究者最后在 i5 CPU 和 Julia 1.3 上完成了测试，结果如表 1 所示非常有竞争力。因为像 PyTorch 那样的成熟深度学习框架，每个 op 的实际开销都至少是 1µs。

这些消失的开销提高了 AD 系统的效率和使用门槛，因此它能以非常细粒度的岔村集成到编程语言中，且还无需担心性能问题。一般而言，自动微分的开销越低，AD 系统的最小可行的内核就越小。因此我们只需要考虑反向传播效率，就能设计出又小又高效的自动微分内核。

其他案例

除了用 Julia 实现 LSTM 模型的自动微分以外，论文还提供了使用可微分编程在其他机器学习领域中的应用示例。

投石机问题——强化学习

图 3： 使用神经网络的替代解决反问题（inverse problem）

基于模型的强化学习对无模型强化学习具有优势，因为高性能的智能体必须能够对其环境动态进行拟合。但是，深度强化学习模型不能够和真实环境结合，限制了它的发展。过去的研究通常使用机器学习框架，可以成功实现现实的物理引擎，但是这种工程的投入非常大，并且和现有的引擎相比有限制，因此在生物学或气象学等领域中的应用较少。

Zygote 可以被用来解决控制问题，能够将模型的反向传播过程控制在一个对梯度的请求中。研究人员使用了投石机问题作为实验案例。他们优化一个可以应对多个目标的神经网络，神经网络的输入是目标的距离和当前的风速，而网络的输出则是投石机的设置（如射程和发射角度）。输入来自一个模拟器，可以作为一个常微分方差，并计算落地点到目标的距离。研究人员然后对比实际结果，并进行反向传播，调整网络权重。研究人员的数据集是随机从一系列目标和风速中选出的。这个智能体可以在笔记本电脑的 CPU 上训练，并可以在常数时间内解决反问题，较直接优化投石机系统的强化学习方法快 100 倍。

计算机视觉

在本案例中，研究人员使用了一个原型生成器，基于期望最终生成的图像，展示优化点光源位置的过程。研究人员定义了一个损失函数，将点光源作为输入，产生图像，并和参考图像对比。按照惯例，梯度可以被提取，并用于更新点光源的位置。

julia> guess = PointLight(Vec3(1.0), 20000.0, Vec3(1.0, 2.0, -7.0))
julia> function loss_function(light)
           rendered_color = raytrace(origin, direction, scene, light, eye_pos)
           rendered_img = process_image(rendered_color, screen_size.w,
                                        screen_size.h)
           return mean((rendered_img .- reference_img).^2)
       end
julia> gs = gradient(x -> loss_function(x, image), guess)

 图 4、5、6： 生成的图片和点光源的关系。 4（左）是最开始生成的图像，5（中）是迭代100 次的结果，6（右）目标图像。