↓ 跳过正文
  1. Posts/

从局部线性化到 RNN:理解 Jacobian 连乘

·4986 字·10 分钟

最近在重拾一些深度学习基础,发现了一个很有意思的视角。

以前在看 RNN(Recurrent Neural Network)的长距离梯度传播公式时,总能看到一堆 Jacobian 矩阵的连乘,比如 \(\frac{\partial h_T}{\partial h_0} = J_T J_{T-1} \dots J_1\)。

如果只是接受这个结论,其实很简单:多维的求导链式法则嘛,把偏导数组装成矩阵乘一乘就是了。但我发现,这种纯代数的理解方式往往在遇到更复杂的梯度消失(vanishing gradient)或爆炸(exploding gradient)现象时,很难建立起直观的物理图像。

这个 Jacobian 到底在描述什么?为什么它的连乘会导致梯度消失或爆炸?这些问题单靠背公式回答不了,得从局部线性化开始重新推一遍。

从导数重新理解“局部线性化”
#

从“斜率”换到局部线性化(Local Linearization)这个视角,更方便接到后面的向量函数上:导数描述输入发生很小变化时,输出如何随之变化。

我们知道,大多数有趣的非线性函数整体上都是弯曲的。但在足够小的局部,它看起来总是像一条直线(或者说切线)。导数,就是用来描述在这个足够小的局部里,输入的变化是如何线性地引起输出的变化的。

从公式上看,对于 \(y = f(x)\),在 \(x_0\) 附近:

$$f(x_0 + \Delta x) \approx f(x_0) + f'(x_0)\Delta x$$

也就是说,当输入有一个微小扰动 \(\Delta x\) 时,输出的扰动可以近似为:

$$\Delta y \approx f'(x_0) \Delta x$$

举个极简的例子,假设 \(f(x) = x^2\)。如果你给 \(x\) 加上一个小的扰动 \(\Delta x\),输出的变化是:

$$(x + \Delta x)^2 = x^2 + 2x\Delta x + (\Delta x)^2$$

当 \(\Delta x\) 极小时,\((\Delta x)^2\) 这个二阶项比 \(\Delta x\) 小得多,可以忽略不计。于是我们得到:

$$\Delta y \approx 2x \Delta x$$

这就把一个二次函数的非线性扰动,转化为了一个只乘上常数 \(2x\) 的线性映射。这就是局部线性化:用简单的比例关系,去近似复杂的非线性变化。

但上面那句"可以忽略不计"其实跳过了一个问题:\((\Delta x)^2\) 在有限大小的 \(\Delta x\) 下并不是零,凭什么说忽略?

看一组数字就能感受到。假设扰动大小为 \(\varepsilon\):

\(\varepsilon\)\(\varepsilon^2\)二阶项 / 一阶项
0.10.010.1
0.010.00010.01
0.0010.0000010.001

随着 \(\varepsilon\) 缩小,\(\varepsilon^2\) 缩小得更快——它相对于 \(\varepsilon\) 的比值本身也在趋向零:

$$\lim_{\varepsilon \to 0} \frac{\varepsilon^2}{\varepsilon} = \lim_{\varepsilon \to 0} \varepsilon = 0$$

这就是所谓"二阶小量":当扰动趋于零时,二阶项消失的速度比一阶项快得多,快到两者的比值也趋于零。

这里有一个我花了一些时间才分清的地方:

  • 在有限的 \(\Delta x\) 下,\((\Delta x)^2\) 不是零。\(\Delta y \approx 2x \Delta x\) 是一个近似,近似就有误差。
  • 但导数 \(f'(x) = 2x\) 不是近似值,它是通过极限精确定义的:
$$f'(x) = \lim_{\Delta x \to 0} \frac{f(x + \Delta x) - f(x)}{\Delta x} = \lim_{\Delta x \to 0} \frac{2x\Delta x + (\Delta x)^2}{\Delta x} = \lim_{\Delta x \to 0} (2x + \Delta x) = 2x$$

极限过程把 \(\Delta x\) 这个残留项彻底消掉了。所以"忽略二阶项"说的是一阶近似有误差,而"导数等于 \(2x\)“说的是极限运算的精确结果——两者是不同层次的事。

多条路径的梯度为什么可以相加
#

上面的 \(x^2\) 例子只有一条路径:\(x \to y\)。但在计算图里,一个变量经常通过多条路径影响最终的 Loss。

我学反向传播的时候,接受了一个规则:“沿路径相乘,跨路径相加”。单路径相乘比较好理解——就是局部线性化一层层传递。但多条路径的梯度为什么可以直接加起来?

从一个具体的例子推一遍。假设:

$$y_0 = 2x, \quad y_1 = 3x$$$$L = y_0 y_1$$

给 \(x\) 加一个扰动 \(\varepsilon\),两个中间变量同时变化:\(\Delta y_0 = 2\varepsilon\),\(\Delta y_1 = 3\varepsilon\)。把 \(\Delta L\) 完整展开:

$$\Delta L = (y_0 + \Delta y_0)(y_1 + \Delta y_1) - y_0 y_1 = y_1 \Delta y_0 + y_0 \Delta y_1 + \Delta y_0 \Delta y_1$$

前两项是一阶项,最后一项 \(\Delta y_0 \Delta y_1\) 是交叉项。代入:

$$\Delta L = (2y_1 + 3y_0)\varepsilon + 6\varepsilon^2$$

按导数的定义取极限:

$$\frac{dL}{dx} = \lim_{\varepsilon \to 0} \frac{\Delta L}{\varepsilon} = \lim_{\varepsilon \to 0} (2y_1 + 3y_0 + 6\varepsilon) = 2y_1 + 3y_0$$

\(6\varepsilon^2\) 除以 \(\varepsilon\) 后变成 \(6\varepsilon\),取极限归零——和前面 \(x^2\) 的道理一样,二阶小量在求导的极限过程中被消除。

写成偏导数的形式:

$$\frac{dL}{dx} = \frac{\partial L}{\partial y_0}\frac{dy_0}{dx} + \frac{\partial L}{\partial y_1}\frac{dy_1}{dx} = y_1 \cdot 2 + y_0 \cdot 3$$

这就是多变量链式法则。每条路径上局部导数相乘,不同路径的贡献相加。

这里有一个需要停下来想想的地方。计算 \(\partial L / \partial y_0\) 时,我们把 \(y_1\) 当作固定常数——可是从计算图来看,\(y_0\) 和 \(y_1\) 都依赖 \(x\)。偏导数问的是另一件事:如果只有 \(y_0\) 单独变化,\(L\) 会怎么变? 每条路径的影响先单独算清楚,最后再汇总——这正是链式法则能拆成求和的原因。

单路径与多路径计算图

不同函数产生的高阶项
#

我当时有一个后续的疑问:上面的交叉项 \(\Delta y_0 \Delta y_1\) 出现在乘法里。如果 \(L\) 不是乘法而是加法,还会出现二阶项吗?如果二阶项和 Loss 的定义相关,链式法则凭什么对所有可微函数都成立?

加法:\(L = y_0 + y_1\),\(\Delta L = \Delta y_0 + \Delta y_1\)。精确等式,没有高阶项。

减法:\(L = y_0 - y_1\),\(\Delta L = \Delta y_0 - \Delta y_1\)。同样精确。链式法则仍然是"跨路径相加”,只不过第二条路径的贡献是负数:\(\partial L / \partial y_1 = -1\)。

乘法:\(L = y_0 y_1\),产生交叉项 \(\Delta y_0 \Delta y_1\)。当两个扰动都是 \(O(\varepsilon)\) 时,交叉项是 \(O(\varepsilon^2)\),在极限中消失。

三次方:\(L = y^3\),展开后有一阶项 \(3y^2 \Delta y\)、二阶项 \(3y(\Delta y)^2\)、三阶项 \((\Delta y)^3\)。高阶项更多,但导数定义的极限只保留一阶。

高阶项的具体形式完全取决于函数本身。有些函数天然就是线性的,不产生高阶项;有些会产生二阶甚至更高阶的项。但链式法则不依赖这些——它只需要函数在当前点可微。

对于一般的 \(L = f(y_0, y_1)\),如果函数在当前点可微:

$$\Delta L = \frac{\partial L}{\partial y_0}\Delta y_0 + \frac{\partial L}{\partial y_1}\Delta y_1 + R$$

余项 \(R = o(\|\Delta \mathbf{y}\|)\),即 \(R\) 相对于扰动的大小趋于零。对具有二阶连续导数的光滑函数可以得到 \(R = O(\|\Delta \mathbf{y}\|^2)\) 的更强结论,但链式法则本身只需要可微这一个条件。

沿路径相乘,跨路径相加,是可微函数局部一阶变化规律的数学结果,不是 Autograd 自行发明的计算规则。

Jacobian:多维函数的局部线性变换
#

到这里可以整理一下概念的层级:一阶导数是一个标量输入对一个标量输出的局部变化率;Gradient 是标量函数对多个输入的偏导数组成的向量;Jacobian 再进一步——多个输出对多个输入的一阶变化关系,组织成矩阵。

一维世界里的局部线性化对应切线。走向多维空间,同样的思想推广成了 Jacobian。

在深度学习里,我们通常面对的是向量到向量的映射,比如 \(\mathbf{y} = f(\mathbf{x})\)。

那么在多维空间中,当我们给输入向量加上一个微小的扰动向量 \(\Delta\mathbf{x}\) 时,输出的扰动 \(\Delta\mathbf{y}\) 会怎样变化?

直接将上面的切线公式推而广之:

$$f(\mathbf{x} + \Delta\mathbf{x}) \approx f(\mathbf{x}) + J\Delta\mathbf{x}$$

也就是说:

$$\Delta\mathbf{y} \approx J\Delta\mathbf{x}$$

很多教材会把 Jacobian 定义成“把所有输出对所有输入的偏导数组装成的一张表”。这个定义没错,但只停留在记账层面。

更直观的理解是:Jacobian 描述了在当前点附近,一个输入的小扰动经过函数以后,会怎样变成输出的小扰动。

可以把它看成局部扰动传播规则,或者当前点附近的局部线性模型。

注意乘法方向:\(\Delta\mathbf{y} \approx J\Delta\mathbf{x}\) 描述的是前向扰动传播——输入端的小扰动经过 \(J\) 映射成输出端的扰动。反向传播时方向反过来:如果已知 Loss 对输出的梯度 \(\mathbf{g}\)(列向量),Loss 对输入的梯度是 \(J^T \mathbf{g}\)。前向乘 \(J\),反向乘 \(J^T\)。

为什么这里自然会出现一个矩阵?因为当我们有多个输入 \((x_1, x_2)\) 变成多个输出 \((y_1, y_2)\) 时,\(x_1\) 的微小变化可能同时影响 \(y_1\) 和 \(y_2\)。我们必须同时描述 \(\frac{\partial y_1}{\partial x_1}, \frac{\partial y_1}{\partial x_2}\) 等等交织在一起的影响。矩阵,就是用来干净利落地把这些局部影响关系组织起来的数学工具。

矩阵本质上是线性变换
#

在 \(\Delta\mathbf{y} \approx J\Delta\mathbf{x}\) 这个公式里,\(J\) 作为一个矩阵,其实扮演了一个线性变换(Linear Transformation)的角色。

\(A\mathbf{x}\) 的操作,直观上就是把一个向量经过特定的拉伸、压缩或旋转,变成另一个向量。

举个非常简单的对角矩阵缩放的例子:

$$A = \begin{bmatrix} 2 & 0 \\ 0 & 3 \end{bmatrix}$$

如果它作用在一个向量上,效果就是把第一个方向放大 2 倍,把第二个方向放大 3 倍。

所以,Jacobian 之所以能够精确描述“局部扰动怎么变化”,正是因为它本身就是一个线性变换算子。在当前这个局部,无论函数本身多么复杂,微小扰动 \(\Delta\mathbf{x}\) 的传播就等价于被这个矩阵 \(J\) 做了拉伸、压缩或旋转等线性操作。

从极简标量 RNN 建立直觉
#

在进阶到真正的 RNN 之前,我们先看一个极简的标量版本 RNN。

假设隐藏状态(hidden state)是一个标量 \(h_t\),并且我们忽略非线性激活函数,只保留核心的时间传播路径:

$$h_t = wh_{t-1} + x_t$$

我们把前几步的时间展开看看:

$$h_1 = wh_0 + x_1$$

$$h_2 = wh_1 + x_2$$

如果我们想知道最初的隐藏状态 \(h_0\) 发生微小变化时,会对 \(h_2\) 产生什么影响,可以通过链式法则求导:

$$\frac{\partial h_2}{\partial h_0} = \frac{\partial h_2}{\partial h_1} \frac{\partial h_1}{\partial h_0} = w \cdot w = w^2$$

推广到 \(T\) 个时间步:

$$\frac{\partial h_T}{\partial h_0} = w^T$$

这个极简的式子直接揭示了循环神经网络在时间维度上的核心问题:

  • 如果 \(|w| \lt 1\),不断相乘以后导数会越来越小,最终趋近于 0,这就导致了最早的扰动无法传导到最后。
  • 如果 \(|w| \gt 1\),不断相乘以后导数会爆炸,模型更新时直接崩溃。

这个因为连乘而导致的指数级衰减/放大的现象,就是 vanishing / exploding gradient 的原型。接下来把同样的推导搬到矩阵 RNN 上。

真正的矩阵 RNN 发生了什么?
#

现在我们来看真实的 RNN 转移方程:

$$h_t = \tanh(W_h h_{t-1} + W_x x_t + b)$$

最核心的区别在于,\(h_t\) 不再是一个标量,而是一个高维向量。这意味着,一个向量状态 \(h_{t-1}\) 如何影响下一个向量状态 \(h_t\),不可能再用一个简单的数字 \(w\) 来描述了。

于是,我们前面提到的“一维导数”自然升级成了 Jacobian 矩阵:

$$J_t = \frac{\partial h_t}{\partial h_{t-1}}$$

回顾前面的局部线性化公式,这就意味着:

$$\Delta h_t \approx J_t \Delta h_{t-1}$$

它回答了这个问题:上一个 hidden state 中的一个小扰动 \(\Delta h_{t-1}\),经过当前时间步的神经网络以后,会怎样被变换传导成新的 hidden state 中的扰动 \(\Delta h_t\)。

在这个过程中,\(W_h\) 负责做隐藏状态间的全局线性映射,而 \(\tanh\) 这种非线性激活函数的导数,则负责在不同分量上对局部扰动进行“自适应的压缩缩放”(比如 \(1-h_t^2\) 项)。它们共同组成了当前这一步独特的局部线性模型 \(J_t\)。

链式法则与 Jacobian 连乘的诞生
#

现在把整条时间线接起来。考虑信息从 \(t=0\) 一直传播到 \(T\):

RNN 状态转移中的 Jacobian

根据多元复合函数的链式法则,\(h_0\) 对 \(h_T\) 的总影响就是把每一步的偏导相乘:

$$\frac{\partial h_T}{\partial h_0} = \frac{\partial h_T}{\partial h_{T-1}} \frac{\partial h_{T-1}}{\partial h_{T-2}} \cdots \frac{\partial h_1}{\partial h_0}$$

代入前面的定义:

$$\frac{\partial h_T}{\partial h_0} = J_T J_{T-1} \cdots J_1$$

这个式子不只是代数上的连乘,它有很直接的传播含义:

  • 每个时间步由一个 \(J_t\) 描述,它回答的是“局部扰动经过这一步会被怎样线性变换”;
  • 跨越几十甚至上百个时间步的梯度传播,就是这些局部线性规则被依次复合。

RNN 的长期梯度传播,本质上就是大量局部线性变换的连续作用。 Jacobian 连乘就是这么来的。

重新审视梯度消失与爆炸
#

有了“矩阵连乘就是局部线性变换连续作用”的视角,我们再回头看梯度消失和爆炸。

在标量时代,是因为 \(w\) 这个数字不断乘自己。而在高维向量时代,一个 Jacobian 矩阵 \(J_t\) 可能在某些方向上放大扰动,在另一些方向上缩小扰动。

当连续成百上千个 Jacobian 矩阵 \(J_T J_{T-1} \dots J_1\) 连乘在一起时:

  • 如果模型在大部分传播路径上持续地把扰动“压扁”,使得最终的复合线性变换把任意微小扰动都缩到了接近零向量的位置,那么 \(\text{gradient} \rightarrow 0\),这就是高维世界里的 vanishing gradient。
  • 如果这条长链条上存在某些方向被持续不断地拉伸放大,使得最终扰动的模长 \(\|\text{gradient}\| \rightarrow \infty\),就出现了 exploding gradient。

结语:从这个视角看 LSTM 和 GRU
#

到这里可以把普通 RNN 的困难说得更准确一点:它并不单纯来自 \(\tanh\) 的导数小于 1,而是 长距离的有效信息和梯度,被迫在每一次状态转移时都经历一整套完整的局部线性变换(Jacobian 连乘)。

步数增加后,复合变换会放大或压缩哪些方向,需要看沿途各个 Jacobian;只看某一步的激活函数导数还不够。

这也自然引出了 LSTM 和 GRU 出现的原因。它们的设计,本质上是在状态传播的路径上引入了更可控的 Gating 机制以及加性路径(additive path)。这使得在很多时间步里,重要信息的 Jacobian 接近于单位矩阵 \(I\)(即扰动原样通过,不放大也不缩小),不需要每一步都被强制重新变换和压缩。

相关文章

Long Short-Term Memory (LSTM) 网络 学习笔记

·761 字·2 分钟
参考资料: # 维基百科_长短期记忆(LSTM) Understanding LSTM Networks [译] 理解 LSTM 网络 LSTM笔记 翻译的比较一般,建议看原文….比如cell还是不要翻译成【细胞】比较好吧…让人以为和生物学的【细胞】有什么关系呢orz