底层码农视角下的自然梯度下降法

虽然很simple，但似乎对于不懂数学的底层码农，如何用我会的数学描述出来已经很艰难了。

我觉得我学不会英语了，我也学不会数学了。

还是补点基础吧

$$\theta_{k+1}=\theta_k-\alpha_k \nabla J(\theta_k)$$

随着距离在某些地方发生变化，梯度下降的方向是能量下降最快的方向并不一定是事实。

比如说我拿KL散度当距离，拿单纯形当参数，一个值越接近零其实反而越精细，微小的参数改动就会导致KL距离大幅度的变化，因为任何两个分布的距离衡量方式变了。这似乎让梯度下降在概率空间中显得不那么自然。

于是自然的想法是，我去修正我的参数的比例关系，使得无论如何我都只尝试从单纯形上去做最速下降，比如单纯性上距离的度量一般就选用KL散度。

KL散度：$D_{KL}(P||Q)=\sum P(x) \log \frac{P(x)}{Q(x)}$，这里 $P,Q$ 是两个带参数的分布。

其实KL散度和熵差不多。

Natural Gradient Descent | Fisher’s Blog

这时候我们的问题就是找一个参数向量的方向 $d$，使得 $J(\theta+d)$ 最小，满足 $KL(p_\theta||p_{\theta+d})=c$，其中 $c$ 是一个常数。

先来看看怎么求解这个问题。

好像不太知道KL散度对参数是怎么影响的，一个直接的想法就是对KL散度做近似。

于是我们对它泰勒展开，注意我们变化是影响两个概率分布的参数，我们把概率分布看成一个点，那么 $P_\theta(x)=p(x\mid \theta), P_{\theta+\delta\theta}(x)=p(x\mid \theta+\delta \theta)$，那么

$$\log P_{\theta+\delta\theta}(x)\approx \log P(x)+\nabla _\theta \log P_\theta(x)\cdot \delta\theta +\frac{1}{2}\delta\theta ^T \mathbf {H_\theta}\log P_\theta(x) \delta\theta$$

其中 $\mathbf H_\theta \log P_\theta(x)$ 就是关于 $\theta$ 的 Hessian矩阵。

score function：我们一般习惯把 $s_\theta(x)=\nabla _\theta \log P_\theta(x)$ 叫做score function，表示似然函数对参数 $\theta$ 的梯度。

$$\nabla_\theta \log P_\theta(x)=\frac{\nabla_\theta P_\theta(x)}{P_\theta(x)}$$

这玩意儿的性质决定了它的期望值是零

$$\int P_\theta(x) \nabla_\theta \log P_\theta(x)\\= \int \nabla_\theta P_\theta(x)=\nabla_\theta \int P_\theta(x)=\nabla_\theta 1=0$$

不知道有没有缺少什么数学严谨性，还有一种理解方式是把那个极限拆开成两个概率分布的差，算一下也是1-1=0。

那么如果第二个参数不变，只变化第一个参数呢？好像不知道会变成什么样子，感觉是大差不差的但是到底差多少不知道。

注意到这时候算全体的期望还是零，所以我们就继续用第二个泰勒展开项。

Fisher Information matrix

对于第二项 $\mathbf H_\theta \log P_\theta(x)$，如果 $\theta$ 是参数向量，写成坐标分量形式就是

$$\mathbf H_\theta \log P_\theta(x)=\left(\frac{\partial \log P_\theta(x)}{\partial \theta_i\partial \theta_j}\right)_{i,j}$$

由于我们还是最后要积分一下要看整体，所以哦我们还是尝试看这个海森矩阵的期望。具体而言，我们尝试去求这个东西逐项的期望。

$$\mathbf F(\theta)=\int P_\theta(x) \mathbf H_\theta \log P_\theta(x) \mathrm dx$$

我们下面也会不加区分地写成 $\mathbf I(\theta)$，写成坐标分量形式就是

$$\mathbf F(\theta)_{ij}=\int P_\theta(x)\frac{\partial \log P_\theta(x)}{\partial \theta_i\partial \theta_j} dx$$

后面这个能怎么化简一下呢？能不能把那个海森矩阵换个简单的形式写一下？

$$\begin{align*} &\quad \mathbf H_\theta \log P_\theta(x)=\nabla _\theta \nabla _\theta \log P_\theta(x)\\ &=\nabla_\theta \frac{\nabla _\theta P_\theta(x)}{P_\theta(x)}\\ &=\frac{P_\theta(x) \nabla_\theta\nabla _\theta P_\theta(x)-(\nabla _\theta P_\theta(x))^2}{P_\theta(x)^2}\\ &=\frac{\nabla_\theta\nabla_\theta P_\theta(x)}{P_\theta(x)}-(\nabla _\theta \log P_\theta(x))^2 \\ \end{align*}$$

注意到如果这时候我们求一个东西的期望

$$\int P_\theta(x) \frac{\nabla_\theta \nabla _\theta P_\theta(x)}{P_\theta(x)} dx\\ =\int \nabla _\theta\nabla _\theta P_\theta(x) dx\\ =\nabla_\theta \nabla _\theta \int P_\theta(x) dx=0$$

所以在期望意义上我们就知道

$$\mathbf F(\theta)=\int P_\theta(x) \mathbf H_\theta \log P_\theta(x) \mathrm dx\\ =-\int P_\theta(x) (\nabla_\theta \log P_\theta(x))^2 \mathrm dx$$

无论怎么表示，我们都把这个求期望后的量的坐标分量 $\mathbf F(\theta)_{ij}$ 定义为 Fisher Information matrix。

KL散度的二阶泰勒展开

综上所述我们得到了KL散度的二阶泰勒展开

$$D_{KL}(P_\theta || P_{\theta+\delta\theta})\approx \frac{1}{2} \delta \theta^T \mathbf F(\theta) \delta \theta$$

值得一提的是，对于另一种对称的推导我们也能得到同样的结论

为了推导KL散度 $KL(P_{\theta+\delta\theta}, P_\theta)$ 的二阶泰勒展开，我们从KL散度的定义出发：

$$KL(P_{\theta+\delta\theta}, P_\theta) = \mathbb{E}_{P_{\theta+\delta\theta}} \left[ \log P_{\theta+\delta\theta}(x) - \log P_\theta(x) \right].$$

首先，对 $\log P_{\theta+\delta\theta}(x)$ 进行泰勒展开：

$$\log P_{\theta+\delta\theta}(x) \approx \log P_\theta(x) + \nabla_\theta \log P_\theta(x) \cdot \delta\theta + \frac{1}{2} \delta\theta^T H_\theta \log P_\theta(x) \delta\theta,$$

其中 $H_\theta$ 是 $\log P_\theta(x)$ 的海森矩阵。

将其代入KL散度的表达式中：

$$KL(P_{\theta+\delta\theta}, P_\theta) \approx \mathbb{E}_{P_{\theta+\delta\theta}} \left[ \nabla_\theta \log P_\theta(x) \cdot \delta\theta + \frac{1}{2} \delta\theta^T H_\theta \log P_\theta(x) \delta\theta \right].$$

接下来，将 $\mathbb{E}_{P_{\theta+\delta\theta}}$ 展开为 $\mathbb{E}_{P_\theta}$ 的函数：

$$\mathbb{E}_{P_{\theta+\delta\theta}} [A] \approx \mathbb{E}_{P_\theta} [A] + \delta\theta^T \mathbb{E}_{P_\theta} [A \nabla_\theta \log P_\theta(x)] + \frac{1}{2} \delta\theta^T \mathbb{E}_{P_\theta} [A H_\theta \log P_\theta(x)] \delta\theta.$$

代入上式并忽略高阶小项，得到：

$$KL(P_{\theta+\delta\theta}, P_\theta) \approx \frac{1}{2} \delta\theta^T I(\theta) \delta\theta,$$

其中 $I(\theta)$ 是Fisher信息矩阵，定义为：

$$I(\theta) = \mathbb{E}_{P_\theta} \left[ \nabla_\theta \log P_\theta(x) \nabla_\theta \log P_\theta(x)^T \right].$$

因此，KL散度 $KL(P_{\theta+\delta\theta}, P_\theta)$ 的二阶泰勒展开为：

$$KL(P_{\theta+\delta\theta}, P_\theta) \approx \frac{1}{2} \delta\theta^T I(\theta) \delta\theta.$$

这个公式的成立基于多元函数的泰勒展开。

$$\mathbb{E}_{P_{\theta+\delta}}[A] \approx \mathbb{E}_{P_{\theta}}[A] + \delta\theta^T \nabla_\theta \mathbb{E}_{P_{\theta}}[A] + \frac{1}{2}\delta\theta^T H_\theta \mathbb{E}_{P_{\theta}}[A] \delta\theta$$

我们可以从概率分布的参数化形式出发，理解其推导过程。

---

FIM是统计模型中自然出现的度量。它提供了参数空间的内在几何结构，反映了模型参数对数据分布的敏感性。

回到之前的问题

找一个参数向量的方向 $d$，使得 $J(\theta+d)$ 最小，满足 $KL(p_\theta||p_{\theta+d})=c$，其中 $c$ 是一个常数。

使用拉格朗日乘子法

$$\min J(\theta+d)+\lambda (D_{KL}(P_\theta||p_{\theta+d})-c)\\ \approx \min_{d}J(\theta)+\nabla_\theta J(\theta)^Td+\frac{1}{2}\lambda d^T\mathbf Fd-\lambda c$$

解得

$$d=-\frac{1}{2\lambda} \mathbf F^{-1} \nabla_\theta J(\theta)$$

---

自然梯度下降 - 知乎

度量

直观上来看，FIM似乎是某种KL散度的近似，但是却有诸如对称等性质，于是我们也尝试把 $\mathbf I$ 看成某种对称的度量。由于这种非常优秀的度量的存在，我们尝试把整个参数空间看成某种流形，每个参数上可以找一个切空间。

我们知道我们定义直接的方向导数是

$$\nabla_v L(p)=\lim_{h\to 0} \frac{L(p+hv)-L(p)}{h}$$

进而定义Riemannian gradient：

$$\nabla _M L(p)= \nabla _v (L(\exp_p(v)))\bigg |_{v=0}$$

进而借由Riemannian exponential map定义Riemannian gradient：

$$p_{t+1}=\exp_{p_t}{(-\alpha_t \nabla_M L(p_t))}$$

这里的 $exp:T_p\to M$ 是把切空间映射到流形上的

但是据说这个计算解需要一个二阶微分方程，直观上是对的，测地线保持内积不变，但是没仔细check过

但是从切空间到流形的映射是困难的，鉴于沿着测地线走比较困难，于是就继续尝试做一阶近似，换言之我用一个新的映射（欧几里德回缩映射） $R_p(v)=p+v$ 来近似 exp。

于是自然梯度下降法就是

$$\theta_{t+1}=\theta_t-\alpha_t ^{NG}\nabla L_\theta(\theta_t)$$

其中自然梯度 ${}^{NG}\nabla$ 定义为

$${}^{NG}\nabla L_\theta(\theta):=g_\theta^{-1} \nabla_\theta L_\theta(\theta)$$

展开为

$$\theta_{t+1}=\theta_t-\alpha_t g_\theta^{-1}(\theta_t)\nabla_\theta L_\theta(\theta_t)$$

可以代入 $g_\theta=\mathbf I$。

之所以这么定义，是因为它意味着梯度下降最快的方向。

为什么自然梯度方向是最速方向

和上面KL散度一样，我们希望 $||d||^2=d^T\mathbf gd\le 1$ 同时最小化 $f(x+\epsilon d)$，就变成了

$$\min_{d}f(\theta)+\nabla_\theta f(\theta)^Td+\frac{1}{2}\lambda d^T\mathbf gd-\lambda c$$

从而解得

$$d=-\frac{1}{2\lambda} \mathbf g^{-1} \nabla_\theta J(\theta)$$

这个观点下，也只是意味着把Fisher Information matrix用作了黎曼度量。之所以用它作为黎曼度量是值得考虑的。

---

另一个空间

既然有了黎曼度量，我们很自然的就能定义两个组参数变化之间的内积。对于在 $\theta$ 处展开的切空间上的两个向量 $\delta\theta_1,\delta\theta_2$，可以定义

$$\left<\delta\theta_1,\delta\theta_2\right>\ =\delta\theta_1^T I(\theta)\delta\theta_2$$

有了内积，我们很自然地想定义两个参数上的联络也就是导数，同样的我们给这组导数施加一些非常简单的条件。

其实我不知道，真不知道，我甚至不知道torsion-free是什么意思，为什么要这样一个setting，但是好像无论如何大家都在用这个东西。我也不知道为什么唯一，但是这好像是个定理的结论。

度量兼容指的是向量沿测地线平移内积不变。

一个严格凸且光滑的生成函数 $F$ 可以构建一个 Hessian 结构，黎曼度量 $g$ 也叫Hessian度量，可以定义为 $F$ 的 Hessian 矩阵。生成函数根据 Legendre-Fenchel对偶变化可以得到对偶函数 $F^*$，也能得到一个 $g^*$ 表示某种对偶 Hessian度量。联络和对偶联络是各自空间的平直联络，Christoffel符号为零。

据说存在唯一的 torsion-free affine connection和度量兼容，这就是Levi-Civita联络 ${}^{LC}\nabla$，所以如果对于一个配备了一组 torsion-free flat connections 流形 $M$，据说 $(\nabla,\nabla^*)$ 满足 $\frac{\nabla+\nabla^*}{2}={}^{LC}\nabla$，这玩意儿我不知道为什么对。

我们之前说到

$${}^{NG}\nabla L_\theta(\theta):=g_\theta^{-1} \nabla_\theta L_\theta(\theta)$$

那么猜想在它的对偶做下降的时候应该是

$$NG^*:\eta_{t+1}=\eta_t-\alpha_t(g^*_\eta)^{-1}(\eta_t)\nabla_\eta L_\theta(\theta(\eta_t)) \\ =\eta_t-\alpha_t(g^*_\eta)^{-1}(\eta_t)\nabla_\eta L_\eta(\eta_t)$$

当然这里 $\eta=\nabla F(\theta)$，$\theta=\nabla F^*(\eta)$，不严谨地可以理解为

$$g_\eta^*(\eta)=\frac{d\theta}{d\eta }=\nabla^2 F^*(\eta)=(\nabla^2F(\theta))^{-1}$$

换一个角度 $g_\theta(\theta)=\nabla^2F(\theta)=\nabla_\theta \nabla_\theta F(\theta)=\nabla_\theta \eta$。

$$\nabla_\theta L_\theta(\theta)=\nabla_\theta (L_\eta(\eta(\theta)))=(\nabla_\theta \eta)(\nabla _\eta L_\eta(\eta))$$

从而我们有

$${}^{NG}\nabla L_\theta(\theta):=g_\theta^{-1}(\theta)\nabla _\theta L_\theta(\theta)\\ =(\nabla_\theta \eta)^{-1} (\nabla _\theta \eta)\nabla _\eta L_\eta(\eta)\\ =\nabla _\eta L_\eta(\eta)$$

结论是在自然梯度下降法等价于对偶参数的常规梯度下降。一个好处是能把由于计算hessian矩阵导致的二阶优化变成一阶优化。

所以现在兜兜转转问题来到了，什么是另一个空间。

example

对于 bernoulli family，直接参数化为选择 $p$ 表示随机变量为1的概率。似然函数 $L(p)=p(1-p)$，对数似然函数 $\log L(p)=\log p-\log (1-p)$。

于是似然函数的二阶梯度是 $\frac{1}{p(1-p)}$，其中一个可以诱导的凸函数取熵函数 $F(p)=p\log p+(1-p)\log (1-p)$, ，定义在 $[0,1]$，$\nabla^2 F^*(p)=\frac{1}{p(1-p)}$

其对偶函数 $F^*(\theta)=\log (1+e^\theta)$，定义在 $\mathbb R$ 上，$\nabla^2 F^*(\theta)=\frac{e^{-\theta}}{(1+e^{-\theta})^2}$。

另外一个例子，如果我是单纯形，参数就是 $\pi$，我取势能函数 $\psi$ 为负的香农熵作为一个势能函数，在这个势能函数意义下。

我要得到 $\sup_{Q} \left<\pi,Q\right>-\pi\log \pi$ ，那么 $Q=\log \pi$，所以我与其在原始空间中做自然梯度下降，不如在 log 空间中做普通的梯度下降。

所以也可以 【凸优化算法】两种视角看待Mirror Descent - 知乎