logo头像

人生之短,取该取舍该舍

K-FAC方法总结

1 摘要

神经网络中的Fisher信息矩阵,既不是对角的,也不是低秩的,并且在某些情况下是完全非稀疏的,因此求逆往往十分复杂。而本文提出的K-FAC是一种可以有效近似神经网络中Fisher信息矩阵的逆的方法,它将Fisher信息矩阵的各个大块(对应于各个层的Fisher信息)近似为两个小得多的矩阵的Kronecker乘积,从而大大提高了求逆的效率。在实际中使用K-FAC进行二阶优化比使用带动量的随机梯度下降速度要快得多,与使用精确Fisher信息逆矩阵的近似自然梯度法/牛顿法的二阶优化方法(如Hessian-free 优化方法)相比,K-FAC在随机优化领域表现很好,因为K-FAC近似曲率逆矩阵的储存耗费和计算量不取决于用于估计它的数据的多少,而只与对角或低秩的近似曲率矩阵相关。K-FAC优化方法是一种能够在神经网络中有效近似自然梯度下降法的方法。

2 神经网络

2.1 Computing Gradient from NN

神经网络中各符号的说明如下表所示:

一个两层神经网络的结构:

即:$s_i=W_i\bar{a}_{i-1},a_i=\phi_i(s_i)$。其中,$\bar{a}=\begin{pmatrix}a\\1\end{pmatrix}$ ,$i=1,2,\cdots,l$,$\phi$为非线性激活函数,各参数的定义见表1。

设损失函数为条件概率密度函数的负对数似然:$L(y,z)=-\log r(y|z)$。则$\forall $参数$v$,定义其梯度为:$\mathcal{D}v=\frac{dL(y,f(x,\theta))}{dv}=-\frac{\log p(y|x,\theta)}{dv}$,$g_i=\mathcal{D}s_i$。定义:$\theta=(vec(W_1)^T\;vec(W_2)^T\;\cdots\;vec(W_l)^T)^T$,则$\theta$的梯度:$\mathcal{D}\theta=(vec(\mathcal{D}W_1)^T\;vec(\mathcal{D}W_2)^T\;\cdots\;vec(\mathcal{D}W_l)^T)^T$ ,计算流程如下:

2.2 Estimating Statistics from NN

一个简单分类网络的前后向传播图如下所示:

由上图可知,对$x\in S_x$,进行前向传播可得到$\bar{a}_i$,并学得$P_{y|x}$,然后进行采样可得到random target $y\sim P_{y|z}$,再利用$y$进行反向传播就可计算得到$g_i$。如:对于二分类问题$y\sim P_{y|x}=\begin{cases}P(x=y_1)=0.72\\P(x=y_2)=0.28\end{cases}$ ,可能采样得到:$y=y_1$。

在4.2节中需计算$F_{i,j}\approx E[\bar{a}_{i-1}\bar{a}_{j-1}^T ]\otimes E[g_i g_j^T]=\bar{A}_{i,j}\otimes G_{i,j}$,从而首先需估计出统计量$\bar{A}_{i,j}$和$G_{i,j}$,于是进行前向和反向传播,然后计算得到:$\bar{A}_{i,j}=E_{\hat{Q}_x}[\bar{a}_i\bar{a}_j^T]=\frac{1}{|S_x|}\displaystyle\sum_{x\in S_x}\bar{a}_i\bar{a}_j^T$,$G_{i,j}=E_{P_{y|x}}[E_{\hat{Q}_x}(g_ig_j^T)]=\frac{1}{|S_{xy}|}\displaystyle\sum_{(x,y)\in S_{xy}}g_ig_j^T$。另外,本文使用指数衰减平均法更新统计量$S$(如$\bar{A}_{i,j},G_{i,j}$),即:$S=(1-\epsilon)S_{old}+\epsilon S_{new},\epsilon=\min\{1-\frac{1}{k},0.95\}$。

3 二阶优化方法

3.1 Natural Gradient Method

Natural Gradient就是使用了Fisher信息的最速下降法,最速下降法的原理如下:

最速下降法

最速下降法通过求解以下优化问题实现:

$\min J(\theta+\Delta\theta)\\subject\;to\;\parallel\Delta\theta\parallel_B<1$

其中:范数$\parallel\cdot\parallel_B$的定义为$\parallel x\parallel_B=(x^TBx)^{\frac{1}{2}}$,它是对向量$x$大小的一种度量。

以上优化问题的解具有以下形式:$\Delta\theta\propto -B^{-1}\nabla_\theta J$。

自然梯度下降法

最速下降法取$\parallel\cdot\parallel_B$为欧几里得范数时,$B=I$,这就是一般的梯度下降法。而当取$\parallel\cdot\parallel_B$取Fisher信息时,$B=F$,其中Fisher信息矩阵$F$为KL散度的二阶近似,这就是自然梯度法。KL散度在度量两个概率分布的差别时与模型的参数化是独立的,而使用欧几里得范数进行度量则差别取决于$\theta$的参数化形式的,下面是一个具体的例子:

自然梯度下降法的更新为:$\theta\gets\theta-\eta F^{-1}\nabla_\theta J$。

3.2 Gauss-Netwon Method

高斯牛顿法和牛顿法的关系

高斯牛顿法是解决非线性最小二乘问题的一种基本方法,它只能处理二次函数。

最小二乘问题:$\min\;f(x)=\frac{1}{2}\displaystyle \sum_{i=1}^mr_i(x)^2=\frac{1}{2}r(x)^T r(x)$ ,其中$r_i(x)$为非线性函数,$r(x)$为函数向量。故$f(x)$的雅克比矩阵为:$J_f(x)=\nabla f(x)=\nabla r(x)r(x)=J_r(x)^T r(x)$,$f(x)$的Hessian矩阵为:$H_f(x)=\nabla^2 f(x)=\nabla r(x)\nabla r(x)^T+r(x)\nabla^2 r(x)=J_r(x)^TJ_r(x)+r(x)H_r(x)\approx J_r(x)^TJ_r(x):=G(x)$,其中$G(x)$为高斯牛顿矩阵。且当$\nabla^2 r(x)=0$时,有$H(x)=G(x)$。牛顿法和高斯牛顿法的更新规则为:

牛顿法:$\Delta=-\frac{\nabla f(x)}{\nabla^2 f(x)}=-H(x)^{-1}J(x),x’=x+\Delta$

高斯牛顿法:$\Delta=-G(x)^{-1}J(x)=-(\nabla r(x)\nabla r(x)^T)^{-1}\nabla r(x)r(x),x’=x+\Delta$

高斯牛顿矩阵和Fisher信息矩阵的关系

当取损失函数$L(y,f(x,\theta))=-\log \;r(y|f(x,\theta))=-\log\;p(y|x,\theta)$时,目标函数:$h(\theta)=E_{\hat{Q}_{x,y}}[L(y,f(x,\theta))]=-\frac{1}{|S|}\displaystyle \underset{(x,y)\in S}{\sum}\log p(y|x,\theta)$,故$h$的Hessian矩阵:$H=E_{\hat{Q}_{x,y}}[\mathcal{D}\theta\mathcal{D}\theta^T]$,而Fisher信息矩阵:$F=E_{P_{x,y}}[\mathcal{D}\theta\mathcal{D}\theta^T]$,其中$P_{x,y}=P_{y|x}(\theta)Q_x=R_{y|f(x,\theta)}Q_x$。

又Martens (2014,arXiv:1412.1193) 证明了:当$R_{y|z}$属于指数分布族时,有$F=G\approx H$,其中$G$为高斯牛顿矩阵。在实验中,我们取$R_{y|z}$为高斯分布,故有$F=G\approx H$。

4 K-FAC方法

4.1 Kronecker product

$\forall A\in R^{mn},B\in R^{pq},X\in R^{qn},u\in R^{n},v\in R^{m},x\in R^{q},y\in R^{p}$,定义:

  • Vectorization:$vec(A_{m,n})=(a_{11},\dots,a_{m1},\dots,a_{1,n},\dots,a_{mn})^T$

  • Kronecker product:

    $A_{m,n}\otimes B_{p,q}=\begin{pmatrix} a_{11}B & \cdots & a_{1n}B\ \vdots & \ddots & \vdots\\a_{m1}B &\cdots & a_{mn}B \end{pmatrix}=M_{mp,nq},v\otimes u=\begin{pmatrix}v_1u\\v_2u\\\vdots \\v_mu\end{pmatrix}=(v_1u_1,\cdots, v_1u_n,\cdots, v_mu_1, \cdots ,v_mu_n)^T$

推论:

  1. $vec(uv^T)=v\otimes u=M_{mn,1}$
  2. $(v\otimes u)(x\otimes y)^T=vx^T\otimes uy^T=M_{mnpq,1}$
  3. $(A_{m,n}\otimes B_{p,q})^{-1}=A_{m,n}^{-1}\otimes B_{p,q}^{-1}=M_{mp,nq}$
  4. $(A_{m,n}\otimes B_{p,q})vec(X_{q,n})=vec(B_{p,q}X_{q,n}A_{m,n}^T)=M_{mp,1}$

4.2 Fisher approximation

Fisher信息矩阵:$F=E[\frac{d\log p(y|x,\theta)}{d\theta}\frac{d\log p(y|x,\theta)}{d\theta}^T]=E[\mathcal{D}\theta\mathcal{D}\theta^T]$

故$F_{i,j}=E[vec(\mathcal{D}W_i)vec(\mathcal{D}W_j)^T]=E[(\bar{a}_{i-1}\otimes g_i)(\bar{a}_{j-1}\otimes g_j)^T]=E[\bar{a}_{i-1}\bar{a}_{j-1}^T \otimes g_i g_j^T]$

对于MINIST(下采样为16*16大小)分类任务,将神经网络中间4层的精确Fisher $F$和使用Kronecker product 近似得到的$\widetilde F$进行对比,并对$\widetilde F$求逆,效果如图2和图3所示:

图2中,左图为$|F|$,中间图为$|\widetilde F|$,右图为$|F|-|\widetilde F|$,其中$|F|$表示对矩阵各元素取绝对值。

图3中,左图为$|\widetilde F|$,中间图为$|\widetilde F^{-1}|$,右图将$|\widetilde F^{-1}|$进行了分块求平均。注:颜色越浅表示值越大,由此可知,Fisher信息逆矩阵中比较大的值集中在对角线附近。

4.3 Approximate Fisher inverse

设$vec(\mathcal{D}\theta)\sim F(0,\Sigma)$,$E[vec(\mathcal{D}\theta)]=0,\Sigma=cov(vec(\mathcal{D}\theta),vec(\mathcal{D}\theta))=E[vec(\mathcal{D}\theta)vec(\mathcal{D}\theta)^T]=F$,记$d_i=vec(\mathcal{D}W_i)$,则$d_i=\displaystyle\sum_{j\neq i}B_{ij}d_j+\epsilon_i$,$vec(\mathcal{D}\theta)=Bvec(\mathcal{D}\theta)+\epsilon$,则$B_{i,i}=0,B_{i,j}=-\frac{(\Sigma^{-1})_{i,j}}{(\Sigma^{-1})_{i,i}}$ ,$D_{i,i}=var(\epsilon_i) =var(d_i|d_j,j\neq i)=\frac{1}{(\Sigma^{-1})_{i,i}}$,故有以下分解:$\Sigma^{-1}=D^{-1}(I-B)$。因此,当$B_{i,j}$很小,即$d_i,d_j$线性相关程度很小时,$(\Sigma^{-1})_{i,j}$很小,从而$(F^{-1})_{i,j}$很小。而$d_1,d_l$的线性相关性是最小的,且有$1=B_{1,1}\geq B_{1,2}\geq\cdots\geq B_{1,l}$,故$F^{-1}$中对角线上的元素是最大的,且其它元素越靠近对角线值越大。故将$\widetilde{F}^{-1}$近似为块对角矩阵和块三角矩阵。

  1. 将$\widetilde{F}^{-1}$近似为块对角矩阵

  1. 将$\widetilde{F}^{-1}$近似为块三角矩阵

设$d_i\sim\mathcal{N}(\Psi_{i,i+1}d_{i+1},\Sigma_{i|i+1}),d_l\sim\mathcal{N}(\overset{\to}0,\Sigma_l)$ ,$d_i=\Psi_{i,i+1}d_{i+1}+\epsilon_i$,故$\Sigma_{i|i+1}=cov(\epsilon_i,\epsilon_i)$。

$cov(d_i,d_{i+1})=cov(\Psi_{i,i+1}d_{i+1}+\epsilon_i,d_{i+1})=\Psi_{i,i+1}cov(d_{i+1},d_{i+1})+0\\\implies E[d_id_{i+1}^T]-E[d_i]E[d_{i+1}]=\Psi_{i,i+1}(E[d_{i+1}d_{i+1}^T]-E^2[d_{i+1}])\\\implies \hat{F}_{i,i+1}-\Psi_{i,i+1}E^2[d_{i+1}]=\Psi_{i,i+1}\hat{F}_{i+1,i+1}-\Psi_{i,i+1}E^2[d_{i+1}]\\\implies \Psi_{i,i+1}=\hat{F}_{i,i+1}\hat{F}_{i+1,i+1}^{-1}$

$cov(d_i,d_{i})=cov(\Psi_{i,i+1}d_{i+1}+\epsilon_i,\Psi_{i,i+1}d_{i+1}+\epsilon_i)=\Psi_{i,i+1}cov(d_{i+1},d_{i+1})\Psi_{i,i+1}^T+\Sigma_{i|i+1}-2\times 0\\\implies E[d_id_{i}^T]-E^2[d_i]=\Psi_{i,i+1}(E[d_{i+1}d_{i+1}^T]-E^2[d_{i+1}])\Psi_{i,i+1}^T+\Sigma_{i|i+1}\\\implies \hat{F}_{i,i}-\Psi_{i,i+1}E^2[d_{i+1}]\Psi_{i,i+1}^T=\Psi_{i,i+1}\hat{F}_{i+1,i+1}\Psi_{i,i+1}^T-\Psi_{i,i+1}E^2[d_{i+1}]\Psi_{i,i+1}^T+\Sigma_{i|i+1}\\\implies \Sigma_{i|i+1}= \hat{F}_{i,i}-\Psi_{i,i+1}\hat{F}_{i+1,i+1}\Psi_{i,i+1}^T$

记$\Psi_{i-1,i}^{\bar{A}}=\bar{A}_{i-1,i}\bar{A}_{i,i}^{-1},\Psi_{i,i+1}^G=G_{i,i+1}G_{i+1,i+1}^{-1}$,则有:

$Ψ_{i,i+1} =\hat F_{i,i+1}\hat F_{i+1,i+1}^{-1} = \widetilde F_{i,i+1} \widetilde F_{i+1,i+1} =(\bar A_{i-1,i }⊗ G_{i,i+1})(\bar A_{i,i} ⊗ G_{i+1,i+1})^{-1} = \Psi_{i-1,i}^{\bar{A}}\otimes \Psi_{i,i+1}^G$

$ \Sigma_{i|i+1}= \hat{F}_{i,i}-\Psi_{i,i+1}\hat{F}_{i+1,i+1}\Psi_{i,i+1}^T=\widetilde{F}_{i,i}-\Psi_{i,i+1}\widetilde{F}_{i+1,i+1}\Psi_{i,i+1}^T\\\qquad\;\;=\bar{A}_{i-1,i-1}\otimes G_{i,i}-\Psi_{i-1,i}^{\bar{A}}\bar{A}_{i,i}\Psi_{i-1,i}^{\bar{A}T}\otimes \Psi_{i,i+1}^GG_{i+1,i+1}\Psi_{i,i+1}^{T}$

根据Grosse and Salakhutdinov (2015),对DGGMs(Pourahmadi, 1999)中$\Sigma^{-1}$的“Cholesky” 分解进行推广,有 :$\hat{F}^{-1}=\Xi^T\Lambda\Xi$,其中

$\Lambda=diag(\Sigma_{1|2}^{-1},\Sigma_{2|3}^{-1},\cdots,\Sigma_{l-1|l}^{-1},\Sigma_{l}^{-1}),\Xi=\begin{pmatrix} I & -\Psi_{1,2} & \; & \; & \; \ \; & I & -\Psi_{2,3}& \; & \; \ \; & \; & I & \ddots & \; \ \; & \; & \; & \ddots & \Psi_{l-1,l}\ \; & \; & \; &\; & I \end{pmatrix}$

由于$\Psi_{i,i+1}=\Psi_{i-1,i}^{\bar{A}}\otimes \Psi_{i,i+1}^G$容易计算,故$\Xi$也容易得到。且若要计算向量$u=\Xi^T v$,则只需计算:$u_i=v_i-\Psi_{i-1,i}^{GT} v_{i-1}\Psi_{i-2,i-1}^{\bar{A}}\;and\;u_1=v_1$。同理,若要计算向量$u=\Xi v$,则只需计算:$u_i=v_i-\Psi_{i,i+1}^{G} v_{i+1}\Psi_{i-1,i}^{\bar{A}T}\;and\;u_l=v_l$。

将以上两种近似Fisher矩阵的方法进行对比,针对前面的MINIST分类任务,对比$\widetilde F、\breve F$与$\widetilde F$的差距,并计算$\widetilde F^{-1}、\breve F^{-1}、\hat F^{-1}$,对比$\widetilde F^{-1}、\breve F^{-1}$与$\widetilde F^{-1}$的差距,效果如图5和图6所示:

图5中,各部分图的表示为:$\begin{pmatrix}|\widetilde F| & |\breve F| & |\widetilde F| - |\breve F| \ |\widetilde F| & |\hat F| & |\widetilde F| - |\hat F| \end{pmatrix} $,图中颜色越深表示值越小,可见$\hat F$保留了大部分的Fisher信息。注:$\hat F=(\hat F^{-1})^{-1}$是块三角矩阵$\hat F^{-1}$的逆,其本身不是块三角矩阵,下图显示了一个对角矩阵C和一个块对角矩阵D的逆的表现形式:

图6中,各部分图的表示为:$\begin{pmatrix}|\widetilde F^{-1}| & |\breve F^{-1} | &| \widetilde F^{-1} | - |\breve F^{-1}| \ |\widetilde F^{-1} | & |\hat F^{-1}| & |\widetilde F^{-1} | - |\hat F^{-1} | \end{pmatrix} $。由右边的图可知,$\widetilde F^{-1}$和$ \hat F^{-1}$的差别已经很小了。

5 更新阻尼(Update damping)

5.1 2nd-order optimization with damping

Martens (2014,arXiv:1412.1193) 证明了:当$R_{y|z}​$属于指数分布族时,有$F=G\approx H​$,其中$G​$为高斯牛顿矩阵。在实验中,我们取$R_{y|z}​$为高斯分布,故有$F=G\approx H​$。于是将目标函数$h(\theta+\delta)​$进行二阶泰勒展开有:$h(\theta+\delta)= \frac{1}{2}\delta^TH\delta+\nabla h(\theta)^T\delta+h(\theta)+\mathcal{O}^2(\delta)​$,故二阶泰勒展开式可近似为:$M(\delta)=\frac{1}{2}\delta^TH\delta+\nabla h(\theta)^T\delta+h(\theta)\approx \frac{1}{2}\delta^TF\delta+\nabla h(\theta)^T\delta+h(\theta)​$,解得:$\delta^*=\underset{\delta}{argmin}\;M(\delta)=-H^{-1}\nabla h(\theta)\approx -F^{-1}\nabla h(\theta)​$。

为防止过拟合,在目标函数中加入正则项$\frac{\eta}{2}\parallel\theta\parallel_2^2$得到:$J(\theta)=h(\theta)+\frac{\eta}{2}\parallel\theta\parallel_2^2$,则$J(\theta)$的Hessian矩阵:$H_J\approx F+\eta I$。为弥补二阶展开误差$\mathcal{O}^2(\delta)$,在$F$中加入$\lambda I$得到更加准确的Hessian矩阵的近似:$F+(\lambda+\eta) I$,即$F+(\lambda+\eta) I$能更好地近似$H$,故:$\delta^*=-(F+(\lambda+\eta) I)^{-1}\nabla h(\theta)$。使用以上两个技巧相当于在近似的$M(\delta)$中加入了$\frac{\lambda+\eta}{2}\parallel\delta\parallel_2^2$,即:$M(\delta)=\frac{1}{2}\delta^T(F+(\lambda +\eta) I)\delta+\nabla h(\theta)^T\delta+h(\theta)$。

5.2 Factored Tikhonov Regularization

令$\stackrel\frown {F}_{i,i}=\widetilde F_{i,i}+(\lambda+\eta)I=\bar{A}_{i-1,i-1}\otimes G_{i,i}+(\lambda+\eta)I\otimes I$,为便于计算$\stackrel\frown {F}_{i,i}^{-1}$,设$\stackrel\frown {F}_{i,i}\approx\stackrel\frown {F}’_{i,i}=(\bar{A}_{i-1,i-1}+\pi_i\sqrt{\lambda+\eta}\;I)\otimes (G_{i,i}+\frac{1}{\pi_i}\sqrt{\lambda+\eta}\;I)$,故有误差项

$error=\stackrel\frown {F}’_{i,i}-\stackrel\frown {F}_{i,i}=\pi_i\sqrt{\lambda+\eta}\;I\otimes G_{i,i}+\frac{1}{\pi_i}\sqrt{\lambda+\eta}\;\bar{A}_{i-1,i-1}\otimes I$

解得:$\pi_i^=\underset{\pi_i}{argmin}\;(F’_{i,i}-F_{i,i})=\sqrt{\frac{tr(\bar{A}_{i-1,i-1})/(d_{i-1}+1)}{tr(G_{i,i})/d_i}}$,其中$d_i$是第$i$层的维度(单元数),于是有:$\stackrel\frown {F}^{-1}_{i,i}\approx (\bar{A}_{i-1,i-1}+\pi_i^\sqrt{\lambda+\eta}\;I)^{-1}\otimes (G_{i,i}+\frac{1}{\pi_i^*}\sqrt{\lambda+\eta}\;I)^{-1}$。

5.3 Re-scaling

对梯度$\Delta$使用学习率$\alpha$进行Re-scaling,从而获得更新的步长:$\delta=\alpha\Delta$。当使用自然梯度法时,记梯度$\Delta=-\stackrel\frown{F}^{-1}\nabla h(\theta)$,其中$\stackrel\frown{F}^{-1}$为使用前面各种技巧所近似得到的Fisher矩阵的逆。从而目标函数$h(\theta+\delta)$的二阶泰勒展开式可近似为:$M(\delta)=M(\alpha\Delta)=\frac{\alpha^2}{2}\Delta^T(F+(\lambda+\eta) I)\Delta+\alpha\nabla h(\theta)^T\Delta+h(\theta)$,其中$F$为在mini-batch上计算得到的准确Fisher矩阵。解得:$\alpha^=\underset{\alpha}{argmin}\;M(\alpha\Delta)=\frac{\nabla h^T\Delta}{\Delta^T(F+(\lambda+\eta) I)\Delta}=\frac{\nabla h^T\Delta}{\Delta^TF\Delta+(\lambda+\eta) \parallel\Delta\parallel_2^2}$,故自然梯度下降法的更新步长为:$\delta=\alpha^\Delta=\frac{\Delta^T\nabla h^T\Delta}{\Delta^TF\Delta+(\lambda+\eta) \parallel\Delta\parallel_2^2}$。

5.4 Momentum

实验发现,在K-FAC中加入动量会有很好的效果。令更新步长为:$\delta=\alpha\Delta+\mu\delta_0$,其中,$\delta_0$为上一次迭代所计算出的更新值,求解$\alpha^=\underset{\alpha}{argmin} \;M(\alpha\Delta+\mu\delta_0)$,$\mu^=\underset{\mu}{argmin} \;M(\alpha\Delta+\mu\delta_0)$,得:

故加入动量的K-FAC方法的更新步长为:$\delta=\alpha^\Delta+\mu^\delta_0$。

实验发现,在随机优化的早期至中期阶段梯度的噪声一般较小,若mini-batch的大小足够大,梯度在后期的噪声也可以很小,此时,加入动量可以实质性地加速梯度下降。效果对比如下图所示:

图7中,横轴表示阻尼系数$\gamma=\sqrt{\lambda+\eta}$,纵轴表示目标函数得到的改进,图中各线的含义如图7右上角所示。

6 调整阻尼系数

6.1 Levenberg-Marquardt Method

对λ使用LM方法进行调参:$\lambda=\begin{cases}w_1\lambda & \rho>3/4 \\\frac{1}{w_1}\lambda & \rho<1/4\end{cases}$ ,其中$\rho=\frac{h(\theta+\delta)-h(\theta)}{M(\delta)-M(0)},0<w_1<1$。此时,当$M\to h$时,$\lambda\to0$,此时趋于NGD方法,这样的好处是:由于此时已到达极值点附近,为防止病态条件下目标函数振荡下降,应考虑曲率,减小更新力度。当$M$不能很好地近似$h$时,$\lambda$增大,此时趋于SGD方法,这样的好处是:由于此时还没有到达极值点附近,无需担心病态条件,故使用SGD方法下降会更快,而若考虑曲率,则会使得更新变缓慢。

6.2 Separate Damping

在计算$\Delta=-\stackrel\frown {F}^{-1}\nabla h$,其中$\stackrel\frown {F}^{-1}_{i,i}\approx (\bar{A}_{i-1,i-1}+\pi_i^\sqrt{\lambda+\eta}\;I)^{-1}\otimes (G_{i,i}+\frac{1}{\pi_i^}\sqrt{\lambda+\eta}\;I)^{-1}$,以及$\alpha^==\frac{\nabla h^T\Delta}{\Delta^TF\Delta+(\lambda+\eta) \parallel\Delta\parallel_2^2}$时,都涉及到了阻尼系数$\lambda+\eta$。这里我们对前者使用不同的阻尼系数,令$\gamma=\sqrt{\lambda+\eta}$,则$\stackrel\frown {F}^{-1}_{i,i}\approx (\bar{A}_{i-1,i-1}+\pi_i^\gamma I)^{-1}\otimes (G_{i,i}+\frac{1}{\pi_i^}\gamma I)^{-1}$,故$\Delta$中的$\gamma$是为了使$\Delta=\underset{\delta}{argmin}\;M(\delta)=\underset{\delta}{argmin}\;(\frac{1}{2}\delta^TH\delta+\nabla h(\theta)^T\delta+h(\theta))\approx-\stackrel\frown {F}^{-1}\nabla h$尽可能准确,而随着迭代的进行,梯度中可能会积累越来越多的误差,因此使用更大的$\gamma$可能会使$\Delta$更准确。而$\alpha^$中的λ,是为了弥补二阶近似误差$h(\theta+\delta)-M(\delta)$,我们希望它随着迭代的进行变得越来越小。因此,对$\Delta$和$\alpha^*$使用不同的阻尼系数$\gamma$和$\lambda+\eta$。

对参数$\gamma$的一个合适的选择是取:$\gamma=\pm M(\alpha^*\Delta)$,且调整$\gamma$的方案为:迭代每隔$T_2$次更新$\gamma$为$\gamma 、w_2\gamma、\frac{1}{w_2}\gamma$中最好的一个。另外,在作者的初始化实验中,使用$\gamma=h(\theta+\delta)$和使用$\gamma=M(\delta)$有类似的效果。

7. 改进计算效率

==TO DO==

8. K-FAC优化方法伪代码

下面是我对K-FAC伪代码的理解:


Algorthm 2 K-FAC方法伪代码

  • 初始化$\theta=\theta_1,k=1$

  • 初始化$\lambda=\lambda_1$(宁愿取很大)

  • 初始化$\gamma=\sqrt{\lambda+\eta}$

  • while $h(\theta_k)>\epsilon_h$ do

    • 选择一个 mini-batch size m
    • 选择一个大小为m的random训练集$S’\subset S$
    • 选择一个大小为$\tau_1m$的子训练集$S_1\subset S’$
    • 选择一个大小为$\tau_2m$的子训练集$S_2\subset S’$
    • 在$S’$上进行前向和后向传播,按照Algorithm 1估计梯度$\nabla h(\theta_k)=\mathcal{D}\theta_k$
    • 采样$y\sim P_{y|x},x\in S_1$,利用$y$再次进行反向传播,得到$g_1^x,\cdots,g_l^x$
    • 计算$\bar{A}_{ij}^{new}=E[\bar a_i \bar a_j^T]=\frac{1}{\tau_1m}\displaystyle\underset{x\in S_1}{\sum}\bar a_i^x \bar (a_j^x)^T$,$G_{ij}^{new}=E[g_i g_j^T]=\frac{1}{\tau_1m}\displaystyle\underset{x\in S_1}{\sum}g^x_i (g^x_j)^T$
    • 更新$\bar A_{ij}=(1-\epsilon)\bar A_{ij}^{new}+\epsilon\bar A_{ij}^{old},G_{ij}=(1-\epsilon)G_{ij}^{new}+\epsilon G_{ij}^{old},\bar A_{ij}^{old}=\bar A_{ij},G_{ij}^{old}=G_{ij}$,其中$\epsilon=\min\{1-1/k,0.95\}$
    • if $(k\mod T_2)=0$
      • $\Gamma=\{\gamma,w_2\gamma,\frac{1}{w_2}\gamma\}$,其中$0<w_2<1$

    end if

    • for each $\gamma \in \Gamma$ do

      • if $(k\mod T_3)=0\;or\;k\leq 3$
        • 更新$\stackrel\frown {F}^{-1}$:$\stackrel\frown {F}^{-1}=(\breve F+\gamma I)^{-1}=diag[(\bar{A}_{i-1,i-1}+\pi_i^\gamma I)^{-1}\otimes (G_{i,i}+\frac{1}{\pi_i^}\gamma I)^{-1}]$ or$\stackrel\frown {F}^{-1}=(\hat F+\gamma I)^{-1}=((\Xi^T\Lambda\Xi)^{-1}+\gamma I)^{-1}$

      end if

      • 计算$\Delta=-\stackrel\frown {F}^{-1}\nabla h(\theta_k)$
      • 更新$\delta\gets\alpha^\Delta+\mu^\delta$,其中$\alpha^,\mu^$见动量部分

    end for

    • 选择$\gamma=\underset{\gamma}{argmin}\;M(\delta)$及相应的$\delta$
    • if $(k\mod T_3)=0$
      • 更新$\lambda$:$\lambda\gets \begin{cases}w_1\lambda & \rho>3/4 \\\frac{1}{w_1}\lambda & \rho<1/4\end{cases}$,其中$\rho=\frac{h(\theta+\delta)-h(\theta)}{M(\delta)-M(0)},0<w_1<1$

    end if

    • $\theta_{k+1}=\theta_k+\delta$
    • $k\gets k+1$

end while