NO.48.tip: 自动微分
背景
以自动方式计算复杂函数的(偏)导数。这里的“复杂”一词是指任意数量的基本操作的组合,例如在深度神经网络中。此任务被称为自动微分(Automatic Differentiation, AD)或 autodiff。自动微分是优化和深度学习的重要组成部分,也用于科学和工程的其他若干领域。
函数形式的微分
在讨论自动微分之前,有必要先复习微分的数学知识。我们将对偏导数使用一种特殊的函数表示法(functional notation),而不是贯穿本书大部分内容的典型表示法。我们暂时将后者称为命名变量表示法(named variable notation)。命名变量表示法依赖于将函数参数与名称相关联。例如,给定一个函数 ,其第一个标量参数的偏导数在点 处,可以记作:
这种表示法引用一个名称 ,隐含或从上下文推断 的参数。另一种表达式是:
其中, 既用作参数的名称(或表达式中的符号),又用作特定的评估点。随着我们将许多函数组合在一起,每个函数可能需要若干个参数,跟踪名称可能会变得越来越复杂。
相反,函数表示法将导数定义为函数上的算子。如果一个函数包含多个参数,则这些参数是通过位置而不是名称来标识的,从而减少了对辅助变量定义的需要。
除了表示法,我们还将依赖一些基本的多变量微积分概念。这包括(偏)导数的概念,函数在一点上的微分或雅可比矩阵,函数作为该点局部线性近似的作用,以及线性映射、矩阵和转置的各种性质。我们将关注有限维的相关设置,并使用 表示 中的标准基。
线性和多线性函数。 我们使用 来表示线性函数 ,并通过 将其应用于 。回想一下,这样的线性映射对应于 中的矩阵,矩阵的列为 ;这两种解释都将被证明是有用的。通常而言,函数组合和矩阵乘法表达式看起来相似:为了将两个线性映射 和 进行组合,可以记作 ,或者,如果不滥用符号,还可以考虑矩阵 。每个线性映射 都对应一个转置 ,这是另一个通过转置对应矩阵来识别的线性映射。
重复使用线性箭头符号,我们可以将其表示为:
一个多线性映射,或者更具体地说是 -线性映射:
其对应于 中的矩阵(或张量)。我们将这种 -线性映射到向量 的应用表示为 。
导数算子
对于开集 和可微函数 ,可将其导数函数(derivative function)表示为:
或者,等价地表示为 。该函数将点 映射到在 处计算的所有偏导数的雅可比矩阵。符号 本身表示导数算子(derivative operator),是一个将函数映射到其导函数的函数。当 时,通过考虑前者的矩阵视图,映射 简化为任意 的标准梯度 。事实上, 符号有时也被描述为算子,因此 是一个函数。当 时,雅可比矩阵是标量值, 是我们熟悉的导数 。
在表达式 中,我们有时会将参数 称为雅可比矩阵的线性化点(linearization point),将 称为扰动(perturbation)。我们将线性化点 和输入扰动 上的映射
称为雅可比-向量积(Jacobian-Vector Product, JVP)。类似地,我们将在线性化点 和输出扰动 上的转置
称为向量-雅可比积(Vector-Jacobian Product, VJP)。
使用映射箭头而不是矩阵可以帮助我们递归地定义高阶导数,接下来我们将展开讨论。
在代码中同时表达这一点,函数与其导数通常是同一个函数;该函数与雅可比矩阵相乘,即 ,而不是明确构造矩阵。更进一步,就这一点而言,我们经常将 同时实现为 ,在程序中编写固定 处的 时,我们通常将其表示为存储器中的数值矩阵。许多代码库为每个(反向)传播输出中的数值矩阵,该函数与雅可比矩阵相乘,即在任何线性化点 和扰动 上。作为标量的简单示例,考虑系统
高阶导数
假设上面的函数 在其定义域 上保持任意可微性。为了再次计算导数,我们可以避免在 时计算 。
记作:
其中, 是表示所有二阶偏导数的双线性映射。例如,在命名变量表示法中,我们可以使用 来表示 。
可以将二阶导数函数 视为两次应用导数算子的结果。也就是说,可以记作 。这种方法可以递归地扩展到任意高阶导数。对于 , 是 -线性映射:
当 时,映射 对应于任意 处的黑塞矩阵。尽管雅可比矩阵和黑塞矩阵足以满足许多机器学习技术,但求解任意高阶导数也并不困难。例如,当将函数展开为泰勒级数近似时,我们可以使用导数算子将其表示为:
多输入
现在考虑带两个参数的函数:
其中,。像 这样的乘积域将有助于对函数的输入组件进行划分, 的导数函数应该是对应于单个输入函数。后者表明,在先前定义的基础上,我们将在两种视角之间切换。多输入往往出现在计算机和程序的上下文中:代码中的许多函数都是为了接受多个参数而编写的,许多基本运算对于多输入,我们可以使用 表示关于第 个参数的导数函数:
如果采用矩阵视角,函数 将一对点 和 映射到 关于第 个参数的所有偏导数的雅可比矩阵:
对于每个线性化点 和扰动 :
采用矩阵视角:
这个约定将简化我们下文讨论的链式法则语句。当 时,两个子矩阵都是标量,并且 简化为偏导数,可以使用命名变量表示法表示为:
然而,表达式 本身具有意义(作为函数),而表达式 在没有进一步上下文的情况下可能是意义模糊的。再次组合算子,便可得到高阶导数。例如,,如果 ,则 在 处的黑塞矩阵为:
组合和扇出
对于某些 和 ,如果 ,则微积分的链式法则表明:
这与多参数函数表示法是如何相互作用的?首先,我们考虑具有扇出的表达式,其中几个子表达式是同一输入的函数。例如,假设两个函数 和 ,并且对于某些函数 ,下式成立:
编写 ,使得 ,式(15)和式(19)表明:
请注意,“+”运算在这里是逐点的意思。从上式还可以看出,如果
换句话说,如果我们书写了多个参数,但没有表现出扇出,那么:
导数的组合和扇出规则使我们能够将复杂的导数计算分解为更简单的导数计算。这就是在处理现代机器学习和数值编程中出现的那种复杂数值计算时,自动微分技术所依赖的技术。
源与流
微分链、电路和程序
自动微分的目的是计算作为输入提供的任意函数的导数。给定一个函数
- 输入扰动 的雅可比-向量积 。
- 输出扰动 的向量-雅可比积 。
换而言之,雅可比-向量积或向量-雅可比积捕获了自动微分的两个基本任务。
我们应该操作哪些函数 作为输入,以及如何表示这些函数,可能是这个设置中最重要的方面。输出组合来描述函数的某种形式化方式。对于基元,我们将在很大程度上考虑自动微分变得越来越具有挑战性。随着语言表达能力的增强,自动微分变得更具挑战性。基元算术运算、约简、收缩、索引、切片和连接等,但我们将如何在很大程序上考虑微分变得容易。链式法则沿着相同的结构分解。链是一类方便的函数表示,因为导数可以划分恰当的链式法则沿着相同的结构分解。
链式组合和链式法则
首先,只取基本操作和链式组合的函数。我们将分阶段加以介绍。
作为一个简单示例,考虑包含三个依次操作的函数 :
根据链式法则,其导数的计算公式为:
该表达式中的括号突出显示了与正向模式自动微分(forward mode automatic differentiation)相对应的从左到右的求值顺序。也就是说,为了计算该雅可比-向量积,需要沿部分因为这些前缀随后立即分别用作后续线性化点:
将该思想扩展到任意的链式组合,可以得到算法 1。
算法 1 链上的正向模式自动微分(雅可比-向量积)
- 输入: 作为链式组合
- 输入:线性化点 和输入扰动
- for do
- 输出:,等于
- 输出:,等于
与之对照,我们可以对式(26)进行转置,以计算输出扰动 的向量-雅可比积:
转置使雅可比映射相对于式(26)中的顺序颠倒过来,现在括号内的计算对应于反向模式自动微分(reverse-mode automatic differentiation)。为了计算该向量-雅可比积,我们可以首先计算原始链前缀 和 ,然后,反向读取这些链前缀作为连续的线性化点:
将该思想扩展到任意的链式组合,可以得到算法 2。
算法 2 链上的反向模式自动微分(向量-雅可比积)
- 输入: 作为链式组合
- 输入:线性化点 和输出扰动
- for do
- for do
- 输出:,等于
- 输出:,等于
尽管链组合强加了一个非常具体的结构,但这种组合方式已经捕获了一些深层神经网络模型,例如多层感知器(假设矩阵乘法是一种原语操作)。
当输出是标量值时(这种情况在深度学习中经常出现,其中输出是损失函数),反向模式自动微分比正向模式自动微分要快。然而,反向模式自动微分在其反向遍历之前存储所有链前缀,因此反向模式自动微分比正向模式消耗更多的内存。在特殊情况下,例如链式操作都是可逆的,研究人员已经提出了一些方法来解决这种内存需求。还可以通过丢弃一些前缀并在需要时重新计算前缀来节省计算的内存。
从链路到电路
当基元可以接受多个输入时,我们可以自然地将链路扩展到电路——基元运算上的有向无环图,有时也被称为计算图。在本节,我们将区分:电路的输入节点(input node),表示函数的参数;基元节点(primitive node),每个基元节点都由基元运算标记。我们假设输入节点没有传入边,(在不失一般性的情况下)每个节点正好有一条传出边,并且图正好有一个汇聚节点。电路的总体功能是从输入节点到汇聚节点的操作组合,其中每个操作的输出根据其输出边输入到其他操作。
自动微分之所以起作用的原因在于,它拥有命名恰当的链式法则,导数沿着链分解。当从链转移到有向无环图时,我们是否需要某种“图式法则”来沿着电路结构分解计算?电路引入了两个新功能:扇入和扇出。在图术语中,扇入只是指传入节点的多条边,扇出是指传出的多条边。
以上内容在函数方面意味着什么呢?当基元操作接受多个参数时,就会发生扇入。在 6.2.1 节中,我们讨论了如何将多个参数视为一个参数,以及如何应用链式法则。扇出需要额外的注意,特别是对于反向模式微分。
答案的要点可以使用一个小例子来说明。考虑图 6-1a 中的电路。运算 在拓扑上先于运算 和运算 ,运算 有两个输出边分别到运算 和运算 。我们可以将 和 分
离以产生两个新电路。第一个电路对应于 ,第二个电路对应于剩余的计算,由下式给出:
我们可以从 和 中恢复完整的函数 ,使用以下辅助函数 dup:
因此, 可以写成链组合:
的电路不包含扇出,像式(23)这样的组合规则表明,其关于 的导数以及其导数的导数,所有这些都是通过链式法则计算而得。同时,将链式法则应用于式(32):
上面的表达式通过从右到左的计算得到 的雅可比-向量积。这一点类似于式(26)中的雅可比-向量积计算,但中间有一个由 dup 函数的雅可比引起的复制操作 。
将 在 上的导数转置:
考虑到从右到左的计算,这也类似于式(28)中的向量-雅可比积计算,但中间有一个求和运算 ,该运算来自 dup 函数的转置雅可比。之所以在这个小例子中使用 dup 运算,更一般地说,是为了在反向模式自动微分中处理扇出,我们可以按照拓扑顺序处理操作——先是正向,然后是反向——然后沿多个输出边对部分向量-雅可比积求和。
算法 3 和算法 4 给出了电路上正向模式微分和反向模式微分的完整描述。为了简洁起见,假设整个电路函数只有一个参数。节点的索引为 。第一个是输入节点,剩余的 个节点通过其操作 来标记。我们将 作为恒等函数。对于每个 ,如果 带 个参数,则设 是其父节点的 个索引的有序列表(由于扇出,可能包含重复索引),设 为其子节点的索引(同样也可能重复)。算法 6.4 采用了更多的约定: 是恒等函数,节点 将 作为其唯一的父节点,节点 1 的子节点是节点 2。
算法 3 正向模式电路微分(雅可比-向量积)
- 输入: 按拓扑顺序组合 ,其中 是恒等函数
- 输入:线性化点 和输入扰动
- for do
- let
- 输出:,等于
- 输出:,等于
算法 4 反向模式电路微分(向量-雅可比积)
- 输入: 按拓扑顺序组合 ,其中 是恒等函数
- 输入:线性化点 和输出扰动
- for do
- let
- for do
- let
- for
- 输出:,等于
- 输出:,等于
扇出是图的特征,但并不是函数的基本特征。通过复制节点,可以从电路表示中删除所有扇出。我们对扇出的兴趣正是为了避免这种情况,从而在算法 6.3 和算法 6.4 中实现高效的表示,进而实现高效的内存使用。
多年来,电路上的反向模式自动微分以各种名称和形式出现。该算法正是神经网络中的反向传播(backpropagation)算法,这是 20 世纪 80 年代引入的一个术语,并分别出现在控制理论和灵敏度的背景下,有关综述,可以参见 Goodfellow、Bengio 和 Courville 的历史笔记。
从电路到程序
图对于引入自动微分算法非常有用,并且图与神经网络应用程序非常吻合。但计算机科学家花了几十年的时间将各种“用于组合表达函数的语言”形式化并进行研究。简单地说,这就是编程语言的用途。对于使用 Python、Haskell 或者 lambda 演算的某些变体表示的数值函数,我们能进行自动微分吗?这些提供了一种更广泛的、直观上更具表现力的描述输入函数的方法 。
在前几节中,由于我们允许更复杂的图结构,因此自动微分方法变得更加复杂。当我们在程序设计语言中引入语法结构时,也会发生类似的情况。我们如何调整自动微分来处理具有循环结构、条件语句和递归调用的语言?并行程序设计结构又如何处理呢?目前,我们部分解决了这些问题,这些问题要求我们更深入地了解编程语言细节,例如类型系统和实现问题。
根据上文中的例子,我们讨论如何处理一个语言构造,即标准 let 表达式。在具有名称或变量绑定方式的语言中,同一变量的多次出现类似于电路中的扇出。我们可以使用函数语言将其记作:
f(x) =
let ax = a(x)
in c(ax, b(ax))
其中, 在绑定后确实出现了两次。
理解语言能力和自动可微性之间的相互作用是计算机科学研究的一个主题。与此同时,函数式语言在最近的自动微分系统中已被证明是非常有效的,无论是广泛使用的还是实验性的。诸如 JAX、Dex 等系统是围绕纯函数编程模型设计的,并且在内部依赖于函数程序表示进行微分。
在 Python 中,该语言所称的“函数”并不总是描述其语法定义中列出的参数的纯函数。“函数”的行为可能依赖于语言所允许的副作用或全局状态。此处,我们特别指的是纯函数的 Python 函数。JAX 的文档详细描述了这种限制。