Intro

自动微分是现代科学计算的一个重要组成部分, 反向传播天然支持大型的神经网络训练, 所以作为兴趣, 稍微了解了一下原理以及简单实现. 参考的材料来自 github 项目 micrograd

自动微分, 顾名思义, 就是机器自动帮我们完成微分操作, 当然, 个人认为这个实际上这个东西并没有听起来这么神, 实现的方法也很简单

对于传统的数值微分, 我们通常存在一些数值的稳定性等问题
对于直接计算微分, 得到的符号的表达式可能特别臃肿

为了平衡这两点, 这就是自动微分的思路: 利用符号表达式连接微分关系, 但我们只储存微分的数值, 具体来说, 自动微分我们需要的所有信息只有微积分学中的 chain rule

(f(g(x)))'=f'(g(x))\cdot g'(x)

也就是说, 我们仍然需要保存符号微分表达式, 但是 chain rule 告诉我们, 计算 $f\circ g$ 的微分, 我们没有必要直接做微分, 我们只需要知道 $f',g'$ 就够了, 剩下就是赋值的操作

所以自动微分的实现, 本质上只是需要我们定义的函数, 或者运算符对象, 需要储存他的导数信息, 这就够了

当然, 仅凭这一点, 并不能完整实现支持神经网络计算需要的反向传播, 在神经网络训练中, 我们通常需要构造一个 loss 函数, 它是巨大数量的参数的函数, 并且复合了超多层(例如深度学习, 实际上就指的是复合的层数). 直接对每个变量计算符号微分表达式是不现实的, 但是 chain rule 告诉我们, 我们只需要一层一层计算就可以了, 具体地, 考察参数 $\mathbf{x}$ 的复合标量函数 $L(f_1,\cdots,f_n)$ , 这里我们显式写出两层, 就有

\frac{\partial L(f(\mathbf{x})) }{\partial x_i}=\sum_{j}\frac{\partial L}{\partial y_j}(f(\mathbf{x}))\cdot \frac{\partial f_j}{\partial x_i}(\mathbf{x})

一个观察是, 仅计算数值, 实际上我们只需要得到 $L$ 的每个导数在给定点的值, 以及 $f_i$ 的导数值, 就可以完成求和得到梯度, 如果 $f_i$ 是更复杂的复合函数, 那就再进行上面的程序继续拆层, 没有区别而最后, 我们总会约化到函数内层的一些简单的微分, 那么先从, $L$ 的导数开始计算, 然后一层一层计算内层的导数, 到变量 $x_i$ 的时候停止, 在计算图上, 计算过程就是从初始点 $L$ 流向参数 $x_i$ , 这个过程就叫反向传播. 相应地, 在计算表达式的值时, 我们称这样的过程为正向传播.

实现一个自动微分, 实际上我们只需要做到:

正确构造带符号微分形式, 以及记录计算图的对象
给出计算图的正确构建

就能给出一个可行版本.

engine: 构造自动微分对象

定义一个 Value , 是反向传播使用的类属性: data(值) , children(子节点), op(算符), grad(梯度)

类中的基本实现:

类属性运算符格式
算符及其反向传播方式
反向传播程序
如果一些交换律, 还需要定义左作用

而对于如何构建反向传播(即在对象的函数表达式中应该储存符号微分), 是下面的流程反向传播计算: 把梯度传递到子节点, 以加法为例

# 反向传播调用后(out = self + other)
self.grad += out.grad
other.grad += out.grad

这是因为

\frac{\partial L}{\partial \text{self}}= \sum_{\text{out}}\frac{\partial L}{\partial \text{out}}\cdot \frac{\partial \text{out}}{\partial \text{self}}

会计算入所有父节点的所有梯度, 而加法的局部偏导数是 1, 于是在反向传播过程中是加法

另一个例子是乘法, 反向传播中具体算法为

out = self * other
self.grad += other.data * out.grad
other.grad += self.data * out.grad

这就是上面的局部 chain rule 的结果

其他函数的实现也是类似的, 每个函数的定义都是如此:

定义函数输出为 Value 类, 包含相同的属性 data,children,op 等
定义反向传播, 并嵌套定义自身的反向传播

反向传播计算图通过构建 topological 排序实现, 做法是通过递归定义不断搜索子节点

实际上计算反向传播的意思是从标量函数的计算图上往回走, 不断累加梯度, 然后在我们给定的参数上停止, 按照上面的反向传播的实现, 我们有自然的要求

计算图停止在每个需要进行梯度下降的参数上
每个需要记录梯度值的参数的父节点必须先被搜索到, 保证梯度传播到方向是越传越深

那么, 从初始节点, 比如 loss 函数, 逐步依靠 Value 的子节点属性向后搜索, 即可构建计算图.

这里我们需要仔细讨论 2. 条件, 这对建图要求很高

拓扑序

这里建图很有技巧, 一种直接的方式是从初始节点直接建图

def build_topo(v):
    if v not in visited:
        visited.add(v)
        topo.append(v)
        for child in v._prev:
            build_topo(child)
build_topo(self)

这样直接建出逆拓扑序的图, 看起来很方便计算反向传播, 按照逆拓扑序一个一个调用即可. 但是仔细想想这种建图方式并不保证我们要求的 2.

考虑这样一个情况, 拓扑图建立的时候遇到了一个节点 a , 但是节点 a 依赖在其他路线上的节点 b , 如果我们运气不好, 先建立了 a 路线的图, 而没有碰到 b 路线, 反向传播中调用了 a 的反向传播就会导致直接炸掉

这是因为我们是从初始节点往深处走一条一条走穿, 再换另一条, 并没有保证遍历顺序, 然而, 只要做上一点改动, 就可以保证这一点:

def build_topo(v):
    if v not in visited:
        visited.add(v)
        for child in v._prev:
            build_topo(child)
        topo.append(v)
build_topo(self)

我们把添加本节点放在添加了所有子节点后, 这导致拓扑序直接被反转, 但排序的性质也会变化:添加本节点必须保证它的子树都被遍历, 那么反过来, 走进一条路必须保证他的前置节点都被调用! 虽然我们的拓扑序被反转了, 但是我们的 2. 要求被满足了, 在使用这个计算图时, 只需要再反转一次计算图, 就可以完成了.

self.grad = 1  # 设置起点梯度, 然后一路反向传播!
for v in reversed(topo):
    v._backward()

至此, 我们就完成了一个最简单的自动微分程序, 自动微分的强大之处在于训练神经网络, 所以我们下面简单介绍一下最简单的神经网络构造

nn: 构造一个简单神经网络

这部分顺便实现了一个神经网络类, 方便实现简单版本的神经网络. 也可以通过这里学习一个神经网络的框架是怎么建立的

首先定义模型(Module)类, 顾名思义, 模型的信息就是它内部的参数, 所以这个类就只包含参数, 然后顺带放了一个重置梯度的函数

在模型的子类中, 我们就可以定义不同的神经网络架构, 这里实现的是一个 Multilayer Perceptron (MLP)

对于 MLP , 我们可以先从一个节点, 一层, 然后一整个 MLP 逐步构建:

由于需要支持反向传播, 在模型中我们使用的所有数学对象都必须是我们前面定义的 Value 类

MLP 的一个神经元储存了所有由前一层指向该神经元的权重, 即系数和偏置, 通过 python 内置的 random 函数进行初始化, 当然我们可以顺便在里面放一个激活函数, 也可以直接使用 Value 类内置的激活函数

定义完一个神经元之后, 就可以定义一个 Layer 定义是简单的, 只需要利用前面定义的单个神经元, 全部拼起来就可以了

可以注意的一点是, 在 layer 定义中的参数, 甚至是后面 MLP 类定义的参数, 都只是递归地按照标量把值放在一个列表中, 这不一定更适配计算机内存的读取. 不过对于这个简单版本, 我们可以不用在意这一部分.

MLP 的定义也只是简单地把 Layer 堆在一起. 不过这里可以注意的一点是看整个架构从 Neuron -> Layer -> MLP 的结构变化:

Neuron: n in , 1 out
Layer: n in , m out
MLP: n in , m out

而每个 n 的意义不同, 对于每一个组件的基本定义就是

初始化: 定义大小, 参数
定义调用类型: 方便后续直接调用进行网络构建
定义参数调用: 方便后续反向传播
如果需要, 可以打印网络类型

训练

有了前面的组件之后, 我们可以简单搭建一个简单神经网络的训练流程

原则上我们还需要写 Dataloader 和 Split 来加载和处理数据, 不过作为最简版本, 我们可以先忽略

训练流程是简单的:

定义模型, 初始化模型
定义 loss 函数和优化方法(比如sgd)
用 loss 函数反向传播, 计算图动态构建(这是一种动态建图的方法), 获取参数的梯度
利用参数的梯度更新参数

即

loss.backward()  # 调用 loss 函数的反向传播, 自动计算相对每个参数的梯度
for p in model.parameters():
    p.data -= lr * p.grad  # 更新参数, 例如这里使用了 sgd 方法

至此, 我们实现了一个最简单版本的神经网络

参考资料

micrograd — Andrej Karpathy 的微型自动微分引擎，本文的主要参考实现
Automatic Differentiation in Machine Learning: a Survey — Baydin et al. 的自动微分综述

最简自动微分

Intro

engine: 构造自动微分对象

拓扑序

nn: 构造一个简单神经网络

训练

参考资料