本文详细介绍了从零开始构建一个支持自动微分的张量(Tensor)库的过程,其核心目标是为后续构建语言模型中的反向传播算法奠定基础。主要内容可分为以下几个关键部分:
1. 张量(Tensor)类的基础设计
文章首先定义了一个基础的Tensor类,用于封装标量数值。该类初始只包含一个value属性。随后,为了支持微分,为Tensor类增加了三个关键属性:
args:记录生成该张量的运算的输入张量。
local_derivatives:记录运算输出对每个输入的局部导数。
derivative:用于存储计算得到的最终梯度。
2. 基本运算与局部导数记录
文章实现了加法(_add)、减法(_sub)和乘法(_mul)三个基本运算函数。这些函数的关键在于,在返回结果张量时,会将输入张量存储到args中,并将预先计算好的局部导数(如加法的导数为1和1,乘法的导数为另一个输入)存储到local_derivatives中,从而构建起计算图的边。
3. 从标量导数到计算图与链式法则
文章从单个运算的标量导数(如 ∂(x+y)/∂x = 1)出发,引出了处理嵌套函数(复合函数)的链式法则。为了可视化与计算复杂的嵌套导数,文章引入了计算图的概念,其中变量用圆圈表示,运算用方框表示,箭头表示数据流动。通过给边标注局部导数,并将从输出到输入的路径上所有边的导数相乘,再对所有路径求和,即可得到最终梯度。
4. 梯度计算算法的实现
文章将上述思想转化为算法。算法使用一个栈来进行广度优先搜索,遍历从最终输出节点(如损失 L)到所有输入叶子节点(如参数 x)的路径。
- 核心过程是:从输出节点开始,携带一个当前梯度(初始为1)。
- 当访问一个中间节点时,将其局部导数与当前梯度相乘,作为新的当前梯度,然后继续向下推送到其输入节点。
- 当到达一个没有
args的叶子节点时,就将累积的梯度累加到该叶子节点的derivative属性上。
这个过程有效地在一次遍历中计算出了输出对所有叶子节点的梯度。
5. 优化与封装
为了代码的简洁性和可用性,文章进行了以下优化:
- 将梯度计算算法封装到
Tensor类的backward方法中。
- 重载了
Tensor类的 __eq__、__add__、__sub__、__mul__ 等特殊方法,使得可以直接使用 +、-、* 和 == 运算符,使代码更直观。
- 利用这些重载后的运算符简化了
backward方法内部的代码。
总结
最终,文章构建了一个能够对仅包含加、减、乘运算的任意数学表达式自动计算导数的轻量级库。其实现原理是通过在运算过程中隐式构建计算图,并在调用backward()方法时,使用基于栈的迭代算法(类似于广度优先搜索)遍历图,应用链式法则累积梯度。这为后续扩展至向量、矩阵乃至神经网络中的反向传播奠定了基础。