Natural Gradient Descent

Natural Gradient Descent

为了了解什么是 natural gradient descent，首先需要知道一些基础知识，包括：

Score function
Fisher Information Matrix
KL 散度及其泰勒展开形式
拉格朗日乘子法

Score function

首先假设我们有一个参数化的模型 $p(x;\theta)$ 。Score function 定义为 log 似然的梯度：
$\begin{aligned} s(\theta) = \nabla_\theta \log p(x;\theta) \end{aligned}$

score function 具有一个性质，也就是其期望为0：
$\begin{aligned} \mathbb{E}_{x \sim p(x;\theta)}[s(\theta)] &= \mathbb{E}_{x \sim p(x;\theta)} [\nabla_\theta \log p(x;\theta)] \\ &= \int p(x;\theta)\frac{\nabla_\theta p(x;\theta)}{p(x;\theta)} \\ &=\nabla \int p(x;\theta) \\ &=0 \end{aligned}$

Fisher Information Matrix

Fisher Information Matrix 定义为 score function 的方差：

$\begin{aligned} F &= \mathbb{E}_{x\sim p(x;\theta)} \Big[\Big(s(\theta)-\mathbb{E}[s(\theta)\Big) \Big(s(\theta)-\mathbb{E}[s(\theta)\Big)^T \Big]\\ &=\mathbb{E}_{x\sim p(x;\theta)} \Big[\nabla_{\theta}p(x;\theta) \nabla_{\theta}p(x;\theta)^T \Big] \end{aligned}$

KL 散度

两个分布 $p(x;\theta)$ 与 $p(x;\theta')$ 之间的KL 散度的定义为：
$KL(p(x;\theta)||p(x;\theta')) = \int p(x;\theta) \log \frac{p(x;\theta)}{p(x;\theta')}$

KL 散度的二阶 Hessian 阵

为了求解自然策略梯度，我们有必要知道 $\text{KL}[p(x;\theta)||p(x;\theta')]$ 的Hessian阵是什么，下面给出KL散度的二阶Hessian阵形式(后面给出具体推导)：
$H_{\theta^{\prime}}\left[\mathrm{KL}\left[p(x ; \theta) \| p\left(x ; \theta^{\prime}\right)\right]\right] = \mathbb{E}_{x \sim p(x ; \theta)}\left[\left(\frac{\nabla_{\theta^{\prime}} p\left(x ; \theta^{\prime}\right)}{p\left(x ; \theta^{\prime}\right)}\right)\left(\frac{\nabla_{\theta^{\prime}} p\left(x ; \theta^{\prime}\right)}{p\left(x ; \theta^{\prime}\right)}\right)^{T}\right]$

然而，KL散度的二阶Hessian阵在 $\theta'=\theta$ 处的取值其实就是 Hessian 阵，也就是说：
$F = H_{\theta^{\prime}}\left[\mathrm{KL}\left[p(x ; \theta) \| p\left(x ; \theta^{\prime}\right)\right]\right]_{ | \theta^{\prime}=\theta}$

推导过程如下：
$\begin{aligned} H_{\theta'}\Big[\text{KL}[p(x;\theta)||p(x;\theta')]\Big] &= H_{\theta'}\Big[\mathbb{E}_{x\sim p(x;\theta)}[\log p(x;\theta)] - \mathbb{E}_{x \sim p(x;\theta)}[\log p(x;\theta')]\Big] \\ &= - H_{\theta'}\Big[\mathbb{E}_{x \sim p(x;\theta)}[\log p(x;\theta')] \Big] \\ &= -\mathbb{E}_{x \sim p(x;\theta)}\Big[H_{\theta'}[\log p(x;\theta')] \Big] \\ &= -\mathbb{E}_{x \sim p(x;\theta)}\Big[ J_{\theta'}[\frac{\nabla_{\theta'}p(x;\theta')}{p(x;\theta')}] \Big] \\ &= -\mathbb{E}_{x \sim p(x;\theta)}\Big[ \frac{H_{\theta'}[p(x;\theta')]p(x;\theta')-\nabla_{\theta'}p(x;\theta')\nabla_{\theta'}p(x;\theta')^T}{p(x;\theta')p(x;\theta')} \Big] \\ &= -\mathbb{E}_{x \sim p(x;\theta)}\Big[ \frac{H_{\theta'}[p(x;\theta')]}{p(x;\theta')}-\Big(\frac{\nabla_{\theta'}p(x;\theta')}{p(x;\theta')}\Big)\Big(\frac{\nabla_{\theta'}p(x;\theta')}{p(x;\theta')}\Big)^T \Big] \\ &= - \int H_{\theta'}[p(x;\theta')] + \mathbb{E}_{x\sim p(x;\theta)}\Big[ \Big(\frac{\nabla_{\theta'}p(x;\theta')}{p(x;\theta')}\Big)\Big(\frac{\nabla_{\theta'}p(x;\theta')}{p(x;\theta')}\Big)^T\Big]\\ &= \underbrace{- H_{\theta'} \Big[ \int p(x;\theta')\Big] }_{0}+ \mathbb{E}_{x\sim p(x;\theta)}\Big[ \Big(\frac{\nabla_{\theta'}p(x;\theta')}{p(x;\theta')}\Big)\Big(\frac{\nabla_{\theta'}p(x;\theta')}{p(x;\theta')}\Big)^T\Big] \\ &=\mathbb{E}_{x\sim p(x;\theta)}\Big[ \Big(\frac{\nabla_{\theta'}p(x;\theta')}{p(x;\theta')}\Big)\Big(\frac{\nabla_{\theta'}p(x;\theta')}{p(x;\theta')}\Big)^T\Big] \end{aligned}$
所以 $\text{KL}[p(x;\theta)||p(x;\theta')]$ 对 $\theta'$ 的导数在 $\theta'=\theta$ 处的取值为Fisher Information Matrix：
$\begin{aligned} H_{\theta'}\Big[\text{KL}[p(x;\theta)||p(x;\theta')]\Big]_{|_{\theta'=\theta}} &=\mathbb{E}_{x\sim p(x;\theta)}\Big[ \Big(\frac{\nabla_{\theta'}p(x;\theta')}{p(x;\theta')}\Big)\Big(\frac{\nabla_{\theta'}p(x;\theta')}{p(x;\theta')}\Big)^T\Big]_{|_{\theta'=\theta}} \\ &=\underbrace{\mathbb{E}_{x\sim p(x;\theta)} \Big[\nabla_{\theta}p(x;\theta) \nabla_{\theta}p(x;\theta)^T \Big]}_{\text{Fisher Information Matrix}} \\ &= F \end{aligned}$

Natural Gradient Descent

基本思路

传统的策略梯度方法选择 $-\nabla_\theta J(\theta)$ 作为参数值 $\theta$ 的更新方向，但是这种方法依赖于欧式空间。而 Natural Gradient Descent 是在分布空间而不是参数空间来更新参数的，其基本思路如下：
$d^*= \arg \min_d L(\theta+d) \quad s.t. \quad KL[p(x;\theta) || p(x;\theta+d)] =c \tag{1}$

根据泰勒二阶
$f(\theta') \approx f(\theta) + (\theta'-\theta)^T\nabla f(\theta) + \frac{1}{2}(\theta'-\theta)^TF(\theta'-\theta)$
可以得到 $\text{KL}[p(x;\theta) || p(x;\theta')]$ 的 $\theta'$ 在 $\theta$ 处的展开式，令 $\theta' -\theta=d$ ，可以得到：
$\begin{aligned} KL[p(x;\theta) || p(x;\theta')] & \approx \underbrace{KL[p(x;\theta)||p(x;\theta)]}_{0} + d^T \nabla_{\theta'}KL[p(x;\theta) || p(x;\theta')]_{|\theta'=\theta} + \frac{1}{2} d^T \underbrace{H_{\theta'}\Big[\text{KL}[p(x;\theta)||p(x;\theta')]\Big]_{|_{\theta'=\theta}}}_{F}d \\ & = 0 + d^T \nabla_{\theta'} \bigg(\mathbb{E}_{x \sim p(x;\theta)} [\log p(x;\theta)] - \mathbb{E}_{x \sim p(x;\theta)} [\log p(x;\theta')] \bigg)_{|\theta' = \theta} + \frac{1}{2} d^T F d \\ & = \frac{1}{2} d^T Fd \end{aligned}$
所以我们通过对(1)中约束进行泰勒展开，得到了下面的优化式子：
$d^* = \arg \min_d L(\theta + d) \quad s.t. \quad \frac{1}{2}d^T Fd = c$
其中对 $L(\theta+d)$ 进行一阶泰勒展开，得到 $L(\theta+d)\approx \mathcal{L}(\theta) + \nabla_\theta \mathcal{L}(\theta)^Td$ 。于是(2)中的优化问题转换成
$d^{*}=\arg \min _{d} \mathcal{L}(\theta) + \nabla_\theta \mathcal{L}(\theta)^Td \quad \text { s.t. } \quad \frac{1}{2} d^{T} F d=c \tag{3}$

可以看出(3)中的优化问题是一个具有等式约束的优化问题，因此我们可以构造拉格朗日函数解上述限制优化问题，首先构造拉格朗日函数 $l(d)$ 为
$l(d) = \mathcal{L}(\theta)+\nabla_\theta \mathcal{L}(\theta)^Td + \lambda (\frac{1}{2} d^{T} F d-c)$

为了求 $L(d)$ 的极小值，令其导数为0，可以得到：
$\begin{aligned} &\text{令}\ \nabla_d l(d) = 0 \\ &\Rightarrow \nabla_\theta \mathcal{L}(\theta) + \lambda F d = 0 \\ & \Rightarrow d = -\frac{1}{\lambda} F^{-1}\nabla_{\theta} \mathcal{L}(\theta) \end{aligned}$
因此，通过对目标函数进行泰勒一次展开，对KL约束进行二次泰勒展开，我们简化了函数的形式，最终得到方向 $d = -\frac{1}{\lambda}F^{-1} \nabla_\theta \mathcal{L}(\theta)$ 。因此参数的更新迭代过程为：
$\theta = \theta - \alpha \frac{1}{\lambda}F^{-1}\nabla_\theta \mathcal{L}(\theta)$

自然策略梯度法的算法如下：

Repeat
1. 进行前向传播计算损失函数 $\mathcal{L}(\theta)$
2. 计算 $\nabla_\theta \mathcal{L}(\theta)$ ，计算Fisher信息矩阵 $F$
3. 计算自然梯度 $d = \frac{1}{\lambda} F^{-1} \nabla_{\theta} \mathcal{L}(\theta)$
4. 更新参数 $\theta = \theta - \alpha d$
Until Convergence