用反向传播(backpropagation)来计算优化目标函数的梯度,是当前机器学习领域的主流方法。近日,牛津与微软等机构的多位学者联合提出一种名为「正向梯度」(forward gradient)的自动微分模式,可以完全抛弃反向传播进行梯度计算。实验证明,在一些问题中,正向梯度的计算时间是反向传播的二分之一。编译 | 张倩编辑 | 陈彩娴反向传播和基于梯度的优化是近年来机器学习(ML)取得重大突破的核心技术。人们普遍认为,机器学习之所以能够快速发展,是因为研究者们使用了第三方框架(如PyTorch、TensorFlow)来解析ML代码。这些框架不仅具有自动微分(AD)功能,还为本地代码提供了基础的计算功能。而ML所依赖的这些软件框架都是围绕 AD 的反向模式所构建的。这主要是因为在ML中,当输入的梯度为海量时,可以通过反向模式的单次评估进行精确有效的评估。自动微分算法分为正向模式和反向模式。但正向模式的特点是只需要对一个函数进行一次正向评估(即没有用到任何反向传播),计算成本明显降低。为此,来自剑桥与微软等机构的研究者们探索这种模式,展示了仅使用正向自动微分也能在一系列机器学习框架上实现稳定的梯度下降。论文地址:https://arxiv.org/pdf/2202.08587v1.pdf他们认为,正向梯度有利于改变经典机器学习训练管道的计算复杂性,减少训练的时间和精力成本,影响机器学习的硬件设计,甚至对大脑中反向传播的生物学合理性产生影响。
1 自动微分的两种模式首先,我们来简要回顾一下自动微分的两种基本模式。
正向模式
给定一个函数 f: θ∈R n,v∈R n,正向模式的AD会计算 f(θ) 和雅可比向量乘积Jf (θ) v,其中Jf (θ) ∈R m×n是f在θ处评估的所有偏导数的雅可比矩阵,v是扰动向量。对于 f : R n → R 的情况,在雅可比向量乘积对应的方向导数用 ∇f(θ)- v表示,即在θ处的梯度∇f对方向向量v的映射,代表沿着该方向的变化率。值得注意的是,正向模式在一次正向运行中同时评估了函数 f 及其雅可比向量乘积 Jf v。此外,获得 Jf v 不需要计算雅可比向量Jf,这一特点被称为无矩阵计算。
反向模式
给定一个函数 f : R n → R m,数值 θ∈R n,v∈R m,AD反向模式会计算f(θ)和雅可比向量乘积v |Jf (θ),其中Jf∈R m×n是f在θ处求值的所有偏导数的雅可比矩阵,v∈R m是一个邻接的矢量。对于f : R n → R和v = 1的情况,反向模式计算梯度,即f对所有n个输入的偏导数∇f(θ)=h ∂f ∂θ1,. . . , ∂f ∂θn i| 。请注意,v |Jf 是在一次前向-后向评估中进行计算的,而不需要计算雅可比Jf 。
运行时间成本
两种AD模式的运行时间以运行正在微分的函数 f 所需时间的恒定倍数为界。反向模式的成本比正向模式高,因为它涉及到数据流的反转,而且需要保留正向过程中所有操作结果的记录,因为在接下来的反向过程中需要这些记录来评估导数。内存和计算成本特征最终取决于AD系统实现的功能,如利用稀疏性。成本可以通过假设基本操作的计算复杂性来分析,如存储、加法、乘法和非线性操作。将评估原始函数 f 所需的时间表示设为 runtime(f),我们可以将正向和反向模式所需的时间分别表示为 Rf×runtime(f) 和 Rb×runtime(f)。在实践中,Rf 通常在1到3之间,Rb通常在5到10之间,不过这些结果都与程序高度相关。
2 方法
正向梯度
定义1给定一个函数 f : R n → R,他们将「正向梯度」 g : R n → R n 定义为:其中,θ∈R n 是评估梯度的关键点,v∈R n 是一个扰动向量,被视为一个多元随机变量v∼p(v),这样 v 的标量分量 vi 是独立的,对所有 i 都有零均值和单位方差,∇f(θ)-v∈R 是 f 在在 v 方向上 θ 点的方向导数。简要地谈一下这个定义的由来。如前所述,正向模式直接给我们提供了方向导数∇f(θ) - v = P i ∂f ∂θi vi,无需计算∇f。将 f 正向评估 n 次,方向向量取为标准基(独热码)向量ei∈R n,i=1 ... n,其中ei表示在第i个坐标上为1、其他地方为0的向量,这时,只用正向模式就可以计算∇f。这样就可以分别评估f对每个输入∂f ∂θi的敏感性,把所有结果合并后就可以得到梯度∇f。为了获得比反向传播更优的运行时间优势,我们需要在每个优化迭代中运行一次正向模式。在一次正向运行中,我们可以将方向v理解为敏感度加权和中的权重向量,即P i ∂f ∂θi vi,尽管这没办法区分每个θi在最终总数中的贡献。因此,我们使用权重向量v将总体敏感度归因于每个单独的参数θi,与每个参数θi的权重vi成正比(例如,权重小的参数在总敏感度中的贡献小,权重大的参数贡献大)。总之,每次评估正向梯度时,我们只需做以下工作:
他们构建了一个正向梯度下降(FGD)算法,用正向梯度g代替标准梯度下降中的梯度∇f(算法1)。在实践中,他们使用小型随机版本,其中 ft 在每次迭代中都会发生变化,因为它会被训练中使用的每一小批数据影响。研究者注意到,算法 1 中的方向导数dt可以为正负数。如果为负数,正向梯度gt的方向会发生逆转,指向预料中的真实梯度。图1显示的两个vk样本,证明了这种行为。在本文中,他们将范围限制在FGD上,单纯研究了这一基础算法,并将其与标准反向传播进行比较,不考虑动量或自适应学习率等其他各种干扰因素。笔者认为,正向梯度算法是可以应用到其他基于梯度算法的优化算法系列中的。