论文整理 Noisy Natural Gradientas Va

2019-12-24 本文已影响0人 WilliamY

1. Fisher Information Matrix （费雪信息矩阵）

假设某概率似然分布 $p(x|\theta)$ ，我们对参数 $\theta$ 求最大似然估计，定义score function（得分函数）如下：
$s(\theta)=\nabla \log p(x|\theta)$
定理: score function对似然分布 $p(x|\theta)$ 求期望为0.
证明:
$\begin{aligned} \mathbb{E}_{p(x | \theta)}[s(\theta)] &=\underset{p(x | \theta)}{\mathbb{E}}[\nabla \log p(x | \theta)] \\ &=\int \nabla \log p(x | \theta) p(x | \theta) \mathrm{d} x \\ &=\int \frac{\nabla p(x | \theta)}{p(x | \theta)} p(x | \theta) \mathrm{d} x \\ &=\int \nabla p(x | \theta) \mathrm{d} x \\ &=\nabla \int p(x | \theta) \mathrm{d} x \\ &=\nabla 1 \\ &=0 \end{aligned}$
为了评估score function，定义Fisher Information Matrix为score function协方差矩阵的期望：
$\mathbf{F}=\underset{p(x | \theta)}{\mathbb{E}}\left[\nabla \log p(x | \theta) \nabla \log p(x | \theta)^{\mathrm{T}}\right]$
实际中我们通过多次采样平均的方法获得其估计值：
$\mathbf{\tilde F}=\frac{1}{N} \sum_{i=1}^N [\nabla \log p(x_i | \theta) \nabla \log p(x_i | \theta)^T]$
定理：Fisher Information Matrix $F$ 是 $p(x|\theta)$ 和 $p(x|\theta')$ 的KL距离在 $\theta'=\theta$ 这一点上的Hessian矩阵。
证明：
$\operatorname{KL}\left[p(x | \theta) \| p\left(x | \theta^{\prime}\right)\right]=\underset{p(x | \theta)}{\mathbb{E}}[\log p(x | \theta)]-\underset{p(x | \theta)}{\mathbb{E}}\left[\log p\left(x | \theta^{\prime}\right)\right]$
对 $\theta'$ 求一阶和二阶导数：
$\begin{aligned} \nabla_{\theta^{\prime}} \mathrm{KL}\left[p(x | \theta) \| p\left(x | \theta^{\prime}\right)\right] &=\nabla_{\theta^{\prime}} \mathbb{E}_{p(x | \theta)}[\log p(x | \theta)]-\nabla_{\theta^{\prime}} \underset{p(x | \theta)}{\mathbb{E}}\left[\log p\left(x | \theta^{\prime}\right)\right] \\ &=-\underset{p(x | \theta)}{\mathbb{E}}\left[\nabla_{\theta^{\prime}} \log p\left(x | \theta^{\prime}\right)\right] \\ &=-\int p(x | \theta) \nabla_{\theta^{\prime}} \log p\left(x | \theta^{\prime}\right) \mathrm{d} x \end{aligned}$
$\nabla_{\theta^{\prime}}^{2} \mathrm{KL}\left[p(x | \theta) \| p\left(x | \theta^{\prime}\right)\right]=-\int p(x | \theta) \nabla_{\theta^{\prime}}^{2} \log p\left(x | \theta^{\prime}\right) \mathrm{d} x$
在 $\theta'=\theta$ 这一点上的Hessian矩阵：
$\begin{aligned} \mathrm{H}_{\mathrm{KL}\left[p(x | \theta) \| p\left(x | \theta^{\prime}\right)\right]} &=-\left.\int p(x | \theta) \nabla_{\theta^{\prime}}^{2} \log p\left(x | \theta^{\prime}\right)\right|_{\theta'=\theta} \mathrm{d} x \\ &=-\int p(x | \theta) \mathrm{H}_{\log p(x | \theta)} \mathrm{d} x \\ &=-\underset{p(x | \theta)}{\mathbb{E}}\left[\mathrm{H}_{\log p(x | \theta)}\right] \\ &=\mathrm{F} \end{aligned}$
【注：Hessian矩阵被定义为实值函数的所有二阶偏导数组成的方块矩阵， $H_{ij}=\frac{\partial ^2 f}{\partial x_i\partial x_j}$ 】

2. Natrual Gradient Desent (自然梯度下降）

对于任意损失函数 $\mathcal L$ ，在取值 $\theta$ 时其梯度最优方向为：
$\frac{-\nabla_{\theta} \mathcal{L}(\theta)}{\left\|\nabla_{\theta} \mathcal{L}(\theta)\right\|}=\lim _{\epsilon \rightarrow 0} \frac{1}{\epsilon} \underset{d \quad s.t. \|d\| \leq \epsilon}{\arg \min } \mathcal{L}(\theta+d)$
定理：当 $d \rightarrow 0$ 时，KL散度的二阶泰勒展开为 $\mathrm{KL}[p(x | \theta) \| p(x | \theta+d)] \approx \frac{1}{2} d^{\mathrm{T}} \mathrm{F} d$
证明：将 $p(x|\theta)$ 简写为 $p_\theta$ ，
$\begin{aligned} \mathrm{KL}\left[p_{\theta} \| p_{\theta+d}\right] & \approx \mathrm{KL}\left[p_{\theta} \| p_{\theta}\right]+\left(\nabla_{\theta'} \mathrm{KL}\left[p_{\theta} \| p_{\theta'}\right] |_{\theta=\theta'}\right)^{\mathrm{T}} d+\frac{1}{2} d^{\mathrm{T}} \mathrm{F} d \\ &=\mathrm{KL}\left[p_{\theta} \| p_{\theta}\right]-\underset{p(x | \theta)}{\mathbb{E}}\left[\nabla_{\theta} \log p(x | \theta)\right]^{\mathrm{T}} d+\frac{1}{2} d^{\mathrm{T}} \mathrm{F} d \end{aligned}$
第一项等于0，第二项由上面的定理 $\underset{p(x | \theta)}{\mathbb{E}}\left[\nabla_{\theta} \log p(x | \theta)\right]=0$ 可知等于0。只剩下第三项。

由该定理，我们得到最优的梯度下降向量d：
$d^*=\underset{d \quad s.t. KL[p_\theta||p_{\theta+d}]=c}{\arg \min } \mathcal{L}(\theta+d)$
其中 $c$ 为常数。把KL设为常数，意思是下降速率恒定不变，而不考虑空间曲率的影响。
用Lagrangian form（拉格朗日法）展开上述优化问题:
$\begin{aligned} d^{*} &=\underset{d}{\arg \min } \mathcal{L}(\theta+d)+\lambda\left(\operatorname{KL}\left[p_{\theta} \| p_{\theta+d}\right]-c\right) \\ & \approx \underset{d}{\arg \min} \mathcal{L}(\theta)+\nabla_{\theta} \mathcal{L}(\theta)^{\mathrm{T}} d+\frac{1}{2} \lambda d^{\mathrm{T}} \mathrm{F} d-\lambda c \end{aligned}$
我们让其对d求导等于0：
$\begin{aligned} 0 &=\frac{\partial}{\partial d} \{\mathcal{L}(\theta)+\nabla_{\theta} \mathcal{L}(\theta)^{\mathrm{T}} d+\frac{1}{2} \lambda d^{\mathrm{T}} \mathrm{F} d-\lambda c \} \\ &=\nabla_{\theta} \mathcal{L}(\theta)+\lambda \mathrm{F} d \\ \lambda \mathrm{F} d &=-\nabla_{\theta} \mathcal{L}(\theta) \\ d &=-\frac{1}{\lambda} \mathrm{F}^{-1} \nabla_{\theta} \mathcal{L}(\theta) \end{aligned}$
如果把 $\frac{1}{\lambda}$ 吸收到学习率中，那么最优的梯度下降向量等于Fisher矩阵的逆乘以损失的梯度。
定义：Natural Gradient（自然梯度）定义为
$\tilde{\nabla}_{\theta} \mathcal{L}(\theta) = \mathrm{F}^{-1} \nabla_{\theta} \mathcal{L}(\theta)$

3.高斯分布的梯度估计

在变分推断中，贝叶斯网络的损失常常被设定为 $ELBO$ ，即后验概率的估计值 $q(\mathbf{w})$ 引发的损失：
$\mathcal{L}_q(\mathbf{w}) = \mathbb{E}_q[\log p(\mathcal{D}|\mathbf{w})]-\lambda D_{KL}[q(\mathbf{w})||p(\mathbf{w})]$
其中 $\mathcal{D}=(x,y)$ ，表示训练数据。
我们令参数集合 $\theta=\{\mu, \Sigma \}$ 分别代表高斯分布的均值和方差， $\bf w \sim \mathcal{N}(\mu, \Sigma)$ 那么其高斯分布的梯度估计为：
$\begin{array}{l}{\nabla_{\boldsymbol{\mu}} \mathbb{E}_{\mathcal{N}(\boldsymbol{\mu}, \mathbf{\Sigma})}[\mathcal{L}(\mathbf{w})]=\mathbb{E}_{\mathcal{N}(\boldsymbol{\mu}, \mathbf{\Sigma})}\left[\nabla_{\mathbf{w}} \mathcal{L}(\mathbf{w})\right]} \\ {\nabla_{\boldsymbol{\Sigma}} \mathbb{E}_{\mathcal{N}(\boldsymbol{\mu}, \mathbf{\Sigma})}[\mathcal{L}(\mathbf{w})]=\mathbb{E}_{\mathcal{N}(\boldsymbol{\mu}, \mathbf{\Sigma})}\left[\nabla_{\mathbf{w}}^{2} \mathcal{L}(\mathbf{w})\right]}\end{array}$
其中 $\bf w$ 表示神经网络的权值矩阵。
这个数学技巧叫做“reprameterization trick（重参数化）”。应用该技巧，我们得到下面的方法。

4. Noisy Natrual Gradient (带噪自然梯度)

$\begin{array}{l}{\tilde{\nabla}_{\mu} \mathcal{L}=\Lambda^{-1} \mathbb{E}_{q}\left[\nabla_{\mathbf{w}} \log p(\mathcal{D} | \mathbf{w})+\lambda \nabla_{\mathbf{w}} \log p(\mathbf{w})\right]} \\ {\tilde{\nabla}_{\Lambda} \mathcal{L}=-\mathbb{E}_{q}\left[\nabla_{\mathbf{w}}^{2} \log p(\mathcal{D} | \mathbf{w})+\lambda \nabla_{\mathbf{w}}^{2} \log p(\mathbf{w})\right]-\lambda \mathbf{\Lambda}}\end{array}$
这里的 $\Lambda=\Sigma^{-1}$ ， $\sim$ 表示自然梯度，。
为了计算上面的式子，对于神经网络的某 $l$ 层考察，设输入为 $a_l \in \mathbb{R} ^ {n_1}$ ，权重矩阵 $W_l \in \mathbb{R}^{n_1 \times n_2}$ ，输出为 $s_l \in \mathbb{R}^{n_2}$ ，忽略偏置项，即 $s_l = W_l^Ta_l$ 。设任意向量 $v$ ，记
$\mathbf{D} v=\nabla_{v} \log p(y | \mathbf{x}, \mathbf{w})$ ，记 $\mathbf{g}_{l}=\mathbf{D} \mathbf{s}_{l}$ ，于是 $\mathbf{D}W_l=a_lg_l^T$ 。
该层的Fisher information matrix
$F_l = \mathbb{E}_{W_l}[vec({\bf D}W_l)vec({\bf D}W_l)^T]=\mathbb{E}[g_lg_l^T \otimes a_la_l^T]$
$\approx \mathbb{E}[g_lg_l^T] \otimes \mathbb{E}[a_la_l^T]=\mathbf{S}_l\otimes \mathbf{A}_l=\tilde{F_l}$
这里有个有趣的观察，从 $\tilde{\nabla}_{\Lambda} \mathcal{L}=0$ ，我们得到不动点：
$\mathbf{\Lambda}=-\mathbb{E}_{q}\left[\frac{1}{\lambda}\nabla_{\mathbf{w}}^{2} \log p(\mathcal{D} | \mathbf{w})+\nabla_{\mathbf{w}}^{2} \log p(\mathbf{w})\right]$
如果 $\lambda =1$ ,则 $\Lambda=H_{-\log p(\mathcal{D},w)}$ ，这和Newton-Raphson更新是一致的。【即牛顿法， $x_{n+1}=x_{n}-\frac{f\left(x_{n}\right)}{f^{\prime}\left(x_{n}\right)}$ 】
简单起见，我们选取 $\mathbf w \sim \mathcal{N}(0,\eta I)$ 为先验，得到 $\nabla_{\bf w} \log p({\bf w})=\eta^{-1} \bf w$ 和 $-\nabla^2_{\bf w} \log p({\bf w})=-\eta^{-1} I$ 。每次迭代，我们随机选择数据 $(x,y)\sim p_{\mathcal D}$ ，采样 ${\bf w} \sim q$ ，
$\begin{array}{l}{\boldsymbol{\mu} \leftarrow \boldsymbol{\mu}+\alpha \mathbf{\Lambda}^{-1}\left[\mathbf{D} \mathbf{w}-\frac{\lambda}{N \eta} \mathbf{w}\right]} \\ {\boldsymbol{\Lambda} \leftarrow\left(1-\frac{\lambda \beta}{N}\right) \boldsymbol{\Lambda}-\beta\left[\nabla_{\mathbf{w}}^{2} \log p(y | \mathbf{x}, \mathbf{w})-\frac{\lambda}{N \eta} \mathbf{I}\right]}\end{array}$
其中 $\alpha$ 和 $\beta$ 是学习率。
进一步简化更新规则， $\Lambda=\frac{N}{\lambda} \overline{\mathbf{F}}+\eta^{-1} \mathbf{I}$
$\overline{\mathbf{F}} \leftarrow(1-\tilde{\beta}) \overline{\mathbf{F}}+\tilde{\beta} \mathcal{D} \mathbf{w} \mathcal{D} \mathbf{w}^{\top}$
$\boldsymbol { \mu } \leftarrow \boldsymbol { \mu } + \tilde { \alpha } \left( \overline { \mathbf { F } } + \frac { \lambda } { N \eta } \mathbf { I } \right) ^ { - 1 } \left[ \mathcal { D } \mathbf { w } - \frac { \lambda } { N \eta } \mathbf { w } \right]$