NO.36.tip： 自然梯度下降

背景

我们讨论自然梯度下降（Natural Gradient Descent, NGD），这是一种用于优化（条件）概率分布 $p_{\theta }(y\mid x)$ 参数的二阶方法。其关键思想是通过测量诱导分布之间的距离来计算参数更新，而不是直接比较参数值。

例如，考虑比较两个高斯分布，$p_{\theta }=p(y\mid \mu ,\sigma )$ 和 $p_{{\theta }^{\prime }}=p(y\mid {\mu }^{\prime },{\sigma }^{\prime })$。参数向量之间的（平方）欧几里得距离分解为 $\Vert \theta -{\theta }^{\prime }{\Vert }^2=(\mu -{\mu }^{\prime }{)}^2+(\sigma -{\sigma }^{\prime }{)}^2$。然而，预测分布的形式为 $\exp \left(-\frac{1}{2{\sigma }^2}(y-\mu {)}^2\right)$，因此需要测量 $\mu$ 相对于 $\sigma$ 的变化。两个单变量高斯分布，其均值相差 $\delta$。两个单变量高斯分布共享相同的小方差 ${\sigma }^2$，，这两个单变量高斯分布共享相同的大方差。很显然，当方差很小时，$\delta$ 的值更重要（就对分布的影响而言）。因此，我们看到这两个参数相互作用，而欧几里得距离无法捕捉到这一点。当我们考虑更复杂的模型（例如深度神经网络）时，这个问题会变得更糟。通过对这种相关性进行建模，自然梯度下降可以比其他梯度方法更快地收敛。

自然梯度下降的关键是根据 KL 散度来测量两个概率分布之间的距离概念。我们可以用 Fisher 信息矩阵来近似。特别地，对于任何给定的输入 $x$，我们有：

$$D_{\mathrm{KL}}(p_{\theta }(y\mid x)\Vert p_{\theta +\delta }(y\mid x))\approx \frac{1}{2}{\delta }^{\top }{F}_x\delta \text{\hspace{1em}(1)}$$

其中，$F_x$ 是 Fisher 信息矩阵。

$$F_x(\theta )=-{\mathbb{E}}_{p_{\theta }(y\mid x)}\left[{\nabla }^2\log p_{\theta }(y\mid x)\right]={\mathbb{E}}_{p_{\theta }(y\mid x)}\left[(\nabla \log p_{\theta }(y\mid x))(\nabla \log p_{\theta }(y\mid x){)}^{\top }\right]\text{\hspace{1em}(2)}$$

我们可以使用 $\frac{1}{2}{\delta }^{\top }F\delta$ 来计算当前分布和更新分布之间的平均 KL 散度，其中 $F$ 是平均 Fisher 信息矩阵：

$$F(\theta )={\mathbb{E}}_{p_{\mathcal{D}}(x)}\left[F_x(\theta )\right]\text{\hspace{1em}(3)}$$

自然梯度下降使用逆 Fisher 信息矩阵作为预处理矩阵，即我们执行以下形式的更新：

$${\theta }_{t+1}={\theta }_t-{\eta }_{t}F({\theta }_t{)}^{-1}{g}_t\text{\hspace{1em}(4)}$$

其中：

$$F^{-1}{g}_t=F^{-1}\nabla \mathcal{L}({\theta }_t)\triangleq \tilde{\nabla }\mathcal{L}({\theta }_t)\text{\hspace{1em}(5)}$$

称为自然梯度。

源与流

自然梯度下降的解释

自然梯度下降作为信赖域方法

我们证明，可以将标准梯度下降解释为优化目标的线性近似，并对参数变化的 ${\ell }_2$ 范数进行惩罚，即如果 ${\theta }_{t+1}={\theta }_t+\delta$，则优化：

$$M_t(\delta )=\mathcal{L}({\theta }_t)+g_t^{\top }\delta +\eta \Vert \delta {\Vert }_2^2\text{\hspace{1em}(6)}$$

现在，让我们将平方距离替换为基于 Fisher 信息矩阵的平方距离，$\Vert \delta {\Vert }_F^2={\delta }^{\top }F\delta$。这相当于白化坐标系（whitened coordinate system）中欧几里得距离的平方 $\varphi =F^{\frac{1}{2}}\theta$，因为：

$$\Vert {\varphi }_{t+1}-{\varphi }_t{\Vert }_2^2=\Vert F^{\frac{1}{2}}({\theta }_t+\delta )-F^{\frac{1}{2}}{\theta }_t{\Vert }_2^2=\Vert F^{\frac{1}{2}}\delta {\Vert }_2^2=\Vert \delta {\Vert }_F^2\text{\hspace{1em}(7)}$$

新目标变成：

$$M_t(\delta )=\mathcal{L}({\theta }_t)+g_t^{\top }\delta +\eta {\delta }^{\top }F\delta \text{\hspace{1em}(8)}$$

求解 ${\nabla }_{\delta }{M}_t(\delta )=0$，得出更新：

$${\delta }_t=-\eta F^{-1}{g}_t\text{\hspace{1em}(9)}$$

这与自然梯度方向相同。因此，我们可以将自然梯度下降视为一种信赖域方法，其中我们对目标使用一阶近似，并在约束中使用 Fisher 信息矩阵距离。

在上面的推导中，我们假设 $F$ 是一个常数矩阵。在大多数问题中，该常数矩阵会在空间中的每个点发生变化，因为我们在一个被称为黎曼流形（Riemannian manifold）的弯曲空间中进行优化。对于某些模型，我们可以有效地计算 Fisher 信息矩阵，即使使用目标的一阶近似，也能够捕获曲率信息。

自然梯度下降作为高斯-牛顿方法

如果 $p(y\mid x,\theta )$ 是一个指数族分布，其自然参数由 $\eta =f(x,\theta )$ 计算，那么我们可以证明自然梯度下降与广义高斯-牛顿（Generalized Gauss-Newton, GGN）方法相同。此外，在在线设置中，这些方法等效于使用扩展卡尔曼滤波器执行顺序贝叶斯推理。

自然梯度下降的优点

使用 Fisher 信息矩阵而不是黑塞矩阵作为预处理矩阵有两个优点。首先，$F$ 总是正定的，而 $H$ 在鞍点上可以有负的特征值，这在高维空间中很普遍。其次，很容易从迷你批次在线近似 $F$，因为 $F$ 是梯度向量的外积的期望值（相对于经验分布）。这与基于黑塞的方法形成了对比，基于黑塞的方法对迷你批次近似引入的噪声更敏感。

此外，与信赖域优化的联系清楚地表明，自然梯度下降以对预测最重要的方式更新参数。这使得该方法能够在参数空间的无信息区域中采取更大的步长，这有助于避免陷入停滞，同时也有助于解决参数高度相关时出现的问题。

我们提出具有不寻常的、高度耦合参数化的二维高斯分布：

$$p(x;\theta )=\frac{1}{2\pi }\exp \left[-\frac{1}{2}{\left(x_1-\left[3{\theta }_1+\frac{1}{3}{\theta }_2\right]\right)}^2-\frac{1}{2}{\left(x_2-\left[\frac{1}{3}{\theta }_1\right]\right)}^2\right]\text{\hspace{1em}(10)}$$

目标是交叉熵损失：

$$\mathcal{L}(\theta )=-{\mathbb{E}}_{p^{\ast }(x)}\left[\log p(x;\theta )\right]\text{\hspace{1em}(11)}$$

该目标的梯度由下式给出：

$${\nabla }_{\theta }\mathcal{L}(\theta )=\left(\begin{array}{c}{\mathbb{E}}_{p^{\ast }(x)}\left[3\left(x_1-\left[3{\theta }_1+\frac{1}{3}{\theta }_2\right]\right)+\frac{1}{3}\left(x_2-\left[\frac{1}{3}{\theta }_1\right]\right)\right]\\ {\mathbb{E}}_{p^{\ast }(x)}\left[\frac{1}{3}\left(x_1-\left[3{\theta }_1+\frac{1}{3}{\theta }_2\right]\right)\right]\end{array}\right)\text{\hspace{1em}(12)}$$

假设 $p^{\ast }(x)=p(x;[0,0])$，于是 Fisher 矩阵是一个常数矩阵：

$$F=\left(\begin{array}{cc}{3}^2+\frac{1}{3^2}& 1\\ 1& \frac{1}{3^2}\end{array}\right)\text{\hspace{1em}(13)}$$

我们将 $\theta$ 空间中的最陡下降与自然梯度法进行了比较。自然梯度法相当于 $\varphi$ 空间中的最大下降。两种方法都从 $\theta =(1,-1)$ 开始。全局最优值为 $\theta =(0,0)$。我们看到，自然梯度方法（蓝点）收敛到这个最优值的速度要快得多，并且采用最短的路径，而最陡的下降则采用非常迂回的路径。我们还看到，白化参数空间中的梯度场呈现为"球形"，这使得下降更加简单和快速。

最后，请注意，无论我们如何对分布进行参数化，自然梯度下降始终保持不变。因此，即使对于高斯的标准参数化，我们也会得到相同的结果。如果我们的概率模型更复杂，例如深度神经网络。

近似自然梯度

自然梯度下降的主要缺点是计算（逆）Fisher 信息矩阵的计算成本。为了加快计算速度，有几种方法对 $F$ 的形式进行了假设，这样就可以有效地求逆。例如，使用对角近似进行神经网络训练，使用低秩加块对角近似，还有假设梯度的协方差可以通过具有低树宽的有向高斯图模型来建模（即 $F$ 的 Cholesky 因子分解是稀疏的）。

有人提出了 KFAC 方法。KFAC 表示"Kronecker 因子近似曲率"（Kronecker factored approximate curvature）。这将深度神经网络的 Fisher 信息矩阵近似为块对角矩阵，其中每个块是两个小矩阵的 Kronecker 乘积。该方法在神经网络的监督学习以及神经策略网络的强化学习上显示出良好的结果。Kronecker 因子近似曲率近似可以使用的平均场分析来证明。此外，有人证明了如果深度神经网络被过度参数化（充当插值器），则 Kronecker 因子近似曲率方法将收敛到深度神经网络的全局最优。

一种更简单的方法是通过使用经验分布代替模型的分布来近似 Fisher 信息矩阵。特别是，定义 $p_{\mathcal{D}}(x,y)=\frac{1}{N}{\sum }_{n=1}^N{\delta }_{x_n}(x){\delta }_{y_n}(y)$，$p_{\mathcal{D}}(x)=\frac{1}{N}{\sum }_{n=1}^N{\delta }_{x_n}(x)$ 和 $p_{\theta }(x,y)=p_{\mathcal{D}}(x)p(y\mid x,\theta )$。然后我们可以按如下公式计算经验 Fisher：

$$F={\mathbb{E}}_{p_{\theta }(x,y)}\left[\nabla \log p(y\mid x,\theta )\nabla \log p(y\mid x,\theta {)}^{\top }\right]\text{\hspace{1em}(14)}$$

$$\approx {\mathbb{E}}_{p_{\mathcal{D}}(x,y)}\left[\nabla \log p(y\mid x,\theta )\nabla \log p(y\mid x,\theta {)}^{\top }\right]\text{\hspace{1em}(15)}$$

$$=\frac{1}{|\mathcal{D}|}\sum _{(x,y)\in \mathcal{D}}\nabla \log p(y\mid x,\theta )\nabla \log p(y\mid x,\theta {)}^{\top }\text{\hspace{1em}(16)}$$

这种近似因计算简单而被广泛使用。特别地，我们可以使用平方梯度向量来计算对角线近似。（这与 AdaGrad 类似，但只使用当前梯度，而不是梯度的移动平均值；后者在执行随机优化时是更好的方法。）

遗憾的是，经验 Fisher 不如真正的 Fisher 有效。其原因是，当我们到达参数空间的平坦部分时，梯度向量为零，经验 Fisher 将变得奇异，因此算法将陷入这个稳定期。然而，真正的 Fisher 将期望值置于输出之上，即该方法将 $y$ 边缘化。如果我们改变参数，这将允许该方法检测输出中的微小变化。这就是为什么自然梯度法比标准梯度法更能"逃离"稳定期的原因。

另一种策略是使用 $F$ 的精确计算，但使用截断共轭梯度方法近似求解 $F^{-1}g$，其中每个共轭梯度步骤使用黑塞向量乘积的有效方法。这被称为无黑塞优化（Hessian free optimization）。然而，这种方法可能很慢，因为计算单个参数更新可能需要多次共轭梯度迭代。

指数族的自然梯度

我们假设 $\mathcal{L}$ 是以下形式的预期损失：

$$\mathcal{L}(\mu )={\mathbb{E}}_{q_{\mu }(z)}\left[\tilde{\mathcal{L}}(z)\right]\text{\hspace{1em}(17)}$$

其中，$q_{\mu }(z)$ 是一个具有矩参数 $\mu$ 的指数族分布。这是变分优化和自然进化策略的基础。

结果表明，矩参数的梯度与自然梯度 $\lambda$ 相同。这遵循链式法则：

$$\frac{d}{d\lambda }\mathcal{L}(\lambda )=\frac{d\mu }{d\lambda }\frac{d}{d\mu }\mathcal{L}(\mu )=F(\lambda ){\nabla }_{\mu }\mathcal{L}(\mu )\text{\hspace{1em}(18)}$$

其中，$\mathcal{L}(\mu )=\mathcal{L}(\lambda (\mu ))$，并使用式 (19) 记作：

$$F(\lambda )={\nabla }_{\lambda }\mu (\lambda )={\nabla }_{\lambda }^{2}A(\lambda )\text{\hspace{1em}(20)}$$

因此：

$${\tilde{\nabla }}_{\lambda }\mathcal{L}(\lambda )=F(\lambda {)}^{-1}{\nabla }_{\lambda }\mathcal{L}(\lambda )={\nabla }_{\mu }\mathcal{L}(\mu )\text{\hspace{1em}(21)}$$

仍然需要计算相对于矩参数的（规则）梯度。具体实现细节将取决于 $q$ 的形式和 $\mathcal{L}(\lambda )$ 的形式。接下来，我们将讨论一些解决该问题的方法。

高斯分布情况下的解析计算

我们假设 $q(z)=\mathcal{N}(z\mid m,V)$。我们现在展示如何解析计算相关的梯度。

根据 2.4.2.5 节，$q$ 的自然参数为：

$${\lambda }^{(1)}=V^{-1}m,{\lambda }^{(2)}=-\frac{1}{2}{V}^{-1}\text{\hspace{1em}(22)}$$

矩参数为：

$${\mu }^{(1)}=m,{\mu }^{(2)}=V+m{m}^{\top }\text{\hspace{1em}(23)}$$

为了简单起见，我们导出标量情况的结果。设 $m={\mu }^{(1)}$，$v={\mu }^{(2)}-({\mu }^{(1)}{)}^2$。通过使用链式法则，矩参数的梯度为：

$$\frac{\partial \mathcal{L}}{\partial {\mu }^{(1)}}=\frac{\partial \mathcal{L}}{\partial m}\frac{\partial m}{\partial {\mu }^{(1)}}+\frac{\partial \mathcal{L}}{\partial v}\frac{\partial v}{\partial {\mu }^{(1)}}=\frac{\partial \mathcal{L}}{\partial m}-2\frac{\partial \mathcal{L}}{\partial v}m\text{\hspace{1em}(24)}$$

$$\frac{\partial \mathcal{L}}{\partial {\mu }^{(2)}}=\frac{\partial \mathcal{L}}{\partial m}\frac{\partial m}{\partial {\mu }^{(2)}}+\frac{\partial \mathcal{L}}{\partial v}\frac{\partial v}{\partial {\mu }^{(2)}}=\frac{\partial \mathcal{L}}{\partial v}\text{\hspace{1em}(25)}$$

仍然需要计算相对于 $m$ 和 $v$ 的导数。如果 $z\sim \mathcal{N}(m,V)$，则根据 Bonnet 定理，我们得到：

$$\frac{\partial }{\partial m_i}\mathbb{E}\left[\tilde{\mathcal{L}}(z)\right]=\mathbb{E}\left[\frac{\partial }{\partial z_i}\tilde{\mathcal{L}}(z)\right]\text{\hspace{1em}(26)}$$

根据 Price 定理，我们得到：

$$\frac{\partial }{\partial V_{ij}}\mathbb{E}\left[\tilde{\mathcal{L}}(z)\right]=c_{ij}\mathbb{E}\left[\frac{{\partial }^2}{\partial z_i\partial z_j}\tilde{\mathcal{L}}(z)\right]\text{\hspace{1em}(27)}$$

其中，当 $i=j$ 时，$c_{ij}=\frac{1}{2}$；否则，$c_{ij}=1$。

在多变量情况下，结果如下所述：

$${\nabla }_{{\mu }^{(1)}}{\mathbb{E}}_{q(z)}\left[\tilde{\mathcal{L}}(z)\right]={\nabla }_m{\mathbb{E}}_{q(z)}\left[\tilde{\mathcal{L}}(z)\right]-2{\nabla }_V{\mathbb{E}}_{q(z)}\left[\tilde{\mathcal{L}}(z)\right]m\text{\hspace{1em}(28)}$$

$$={\mathbb{E}}_{q(z)}\left[{\nabla }_z\tilde{\mathcal{L}}(z)\right]-{\mathbb{E}}_{q(z)}\left[{\nabla }_z^2\tilde{\mathcal{L}}(z)\right]m\text{\hspace{1em}(29)}$$

$${\nabla }_{{\mu }^{(2)}}{\mathbb{E}}_{q(z)}\left[\tilde{\mathcal{L}}(z)\right]={\nabla }_V{\mathbb{E}}_{q(z)}\left[\tilde{\mathcal{L}}(z)\right]\text{\hspace{1em}(30)}$$

$$=\frac{1}{2}{\mathbb{E}}_{q(z)}\left[{\nabla }_z^2\tilde{\mathcal{L}}(z)\right]\text{\hspace{1em}(31)}$$

因此，我们看到自然梯度依赖于损失函数 $\tilde{\mathcal{L}}(z)$ 的梯度和黑塞矩阵。

一般情况下的随机近似

一般来说，很难采用解析方式计算自然梯度。然而，我们可以计算蒙特卡罗近似。为了理解这一点，我们假设 $\mathcal{L}$ 是以下形式的预期损失：

$$\mathcal{L}(\mu )={\mathbb{E}}_{q_{\mu }(z)}\left[\tilde{\mathcal{L}}(z)\right]\text{\hspace{1em}(32)}$$

根据式 (21)，自然梯度由下式给出：

$${\nabla }_{\mu }\mathcal{L}(\mu )=F(\lambda {)}^{-1}{\nabla }_{\lambda }\mathcal{L}(\lambda )\text{\hspace{1em}(33)}$$

对于指数族分布，上述公式右侧的两项都可以写成期望值，因此可以通过蒙特卡罗方法进行近似。为了理解这一点，请注意：

$$F(\lambda )={\nabla }_{\lambda }\mu (\lambda )={\nabla }_{\lambda }{\mathbb{E}}_{q_{\lambda }(z)}\left[T(z)\right]\text{\hspace{1em}(34)}$$

$${\nabla }_{\lambda }\mathcal{L}(\lambda )={\nabla }_{\lambda }{\mathbb{E}}_{q_{\lambda }(z)}\left[\tilde{\mathcal{L}}(z)\right]\text{\hspace{1em}(35)}$$

如果 $q$ 是可重新参数化的，那么我们可以应用重新参数化技巧将梯度移动到期望算子内部。这样我们就从 $q$ 中采样出 $z$，然后计算梯度并求平均值，之后可以将得到的随机梯度传递给随机梯度下降。

熵函数的自然梯度

我们将讨论如何计算指数族分布的熵的自然梯度，这在进行变分推理时很有用。自然梯度的计算公式为：

$${\tilde{\nabla }}_{\lambda }\mathbb{H}(\lambda )=-{\nabla }_{\mu }{\mathbb{E}}_{q_{\mu }(z)}\left[\log q(z)\right]\text{\hspace{1em}(36)}$$

其中，根据式 (37)，我们得到：

$$\log q(z)=\log h(z)+T(z{)}^{\top }\lambda -A(\lambda )\text{\hspace{1em}(38)}$$

由于 $\mathbb{E}[T(z)]=\mu$，因此：

$${\nabla }_{\mu }{\mathbb{E}}_{q_{\mu }(z)}\left[\log q(z)\right]={\nabla }_{\mu }{\mathbb{E}}_{q(z)}\left[\log h(z)\right]+{\nabla }_{\mu }{\mu }^{\top }\lambda (\mu )-{\nabla }_{\mu }A(\lambda )\text{\hspace{1em}(39)}$$

其中，$h(z)$ 是基础度量值。由于 $\lambda$ 是 $\mu$ 的函数，我们有：

$${\nabla }_{\mu }{\mu }^{\top }\lambda =\lambda +({\nabla }_{\mu }\lambda {)}^{\top }\mu =\lambda +(F_{\lambda }^{-1}{\nabla }_{\lambda }\lambda {)}^{\top }\mu =\lambda +F_{\lambda }^{-1}\mu \text{\hspace{1em}(40)}$$

由于 $\mu ={\nabla }_{\lambda }A(\lambda )$，我们有：

$${\nabla }_{\mu }A(\lambda )=F_{\lambda }^{-1}{\nabla }_{\lambda }A(\lambda )=F_{\lambda }^{-1}\mu \text{\hspace{1em}(41)}$$

因此：

$$-{\nabla }_{\mu }{\mathbb{E}}_{q_{\mu }(z)}\left[\log q(z)\right]=-{\nabla }_{\mu }{\mathbb{E}}_{q(z)}\left[\log h(z)\right]-\lambda \text{\hspace{1em}(42)}$$

如果我们假设 $h(z)=\text{常数}$，也就是通常的情况，那么得到：

$${\tilde{\nabla }}_{\lambda }\mathbb{H}(\lambda )=-\lambda \text{\hspace{1em}(43)}$$