Lorn3's blog

Hessian矩阵系列串讲

定义

假设有一实值函数f(x1,x2,,xn),若f的所有二阶偏导数都存在且在定义域里连续,那么我们定义函数f的Hessian矩阵为如下一个n×n的方阵:

𝐇=[2fx122fx1x22fx1xn2fx2x12fx222fx2xn2fxnx12fxnx22fxn2]

或使用下标标记表示为:

𝐇ij=2fxixj

一般性质

对称性:

对于机器学习中遇到的大多数函数(特别是那些具有连续二阶偏导数的函数),混合偏导数的求导顺序无关紧要。这被称为克莱罗定理或施瓦茨定理,它说明了混合偏导数的相等性:

2fxixj=2fxjxi

这意味着Hessian矩阵是对称的,即𝐇=𝐇

正定性与局部最优值的相关性

通过Hessian矩阵的正定性,我们可以判断该点是局部最小值,局部最大值还是鞍点

Hessian在神经网络中的计算

Hessian矩阵在神经网络计算的许多方面有着重要作用,包括:

对于Hessian矩阵的众多应用而言,一个重要的需要考虑的问题是计算效率。在神经网络中有W个参数(包括权值和偏置),那么Hessian矩阵的大小就是W×W,那么计算Hessian矩阵的计算量为O(W2)。这在具有大量参数的神经网络中是难以接受的,因此我们需要进行一些高效的近似。

对角近似

这个方法最早由Yan LeCun在OBD剪枝方法中提出

我们只保留Hessian矩阵的对角线元素Hi,i,把非对角线元素置为0.(这个方法在很大程度可以可以方便求逆)

对于神经网络中的第j个神经元的输入aj对应权重为wji:

2Enwji2=2Enaj2zi2

其中zi是上一层神经元的输出。

2Enaj2可以通过链式法则递归计算(类似于反向传播):

2Enaj2=aj[h(aj)kwkjEnak]=h(aj)2k,kwkjwkj2Enwkwk+h(aj)kwkjEnan

忽略二阶导中的非对角线项kk:

2Enaj2h(aj)2kwkj22Enak2+h(aj)kwkjEnak

从而一次反向传播便可以计算出来,时间复杂度为O(W)

但是如你所见,在这个方法下,Hessian矩阵完全退化为了一个diag,这在神经网络中是不合理的,因为非线性层的引入会让Hessian矩阵的交叉项不为0,更重要的是这种误差会随着网络层数的堆砌不断放大。

因此针对这个问题在对角近似上近年来不断有相关文章发表,下面对我阅读过的一些进行简单介绍:

  1. HesScale:在LeCun的方法的骨架上,针对常见网络的最后一层的softmax+CE结构可以通过一个简单的公式求出diag(H)的精确值:
2En2aj2=ppp

其中p为softmax(aj),表示向量的逐元素乘积,为方便书写我们将{a1,,an}写作z

证明如下:

​ 预测𝐩=softmax(z),即pi=exp(zi)j=1kexp(zj),设目标分布为tΔK1,iti=1,对于单样本而言,其损失可以写作:

(z)=j=1Ktilogpi=tz+logj=1Kexp(zj)

求一阶导我们有:

z=t+zlogj=1Kexpzj=t+𝐩

求二阶导相当于𝐩对z求导。我们知道softmax的Jocobian为

pizj=pi(δijpj),Jsoftmax(z)=diag(𝐩)𝐩𝐩

其中δi,j为Kronecker函数,写成矩阵形式,则只有对角线元素为1,其余为0.

因此我们有:

z2=diag(𝐩)pp

于是我们可以得到:

2L2aj=pj(1pj)
  1. AdaHessian:

    具体而言,AdaHessian在进行Hessian矩阵的对角近似的时候利用了Hutchinson估计,它通常也被用来对矩阵的迹进行估计:

    ​ 对任意矩阵𝐀d×d,若随机向量z=(z1,,zd)满足:

𝔼[zi]=0𝔼[zizj]=δij

​ 那么有:

𝔼[z(Az)]=diag(A)

​ 下面进行证明:

​ 记估计量d=z(Az)的第i个分量:

di=zi(Az)i=zij=1dAijzj=j=1dAijzizj𝔼[di]=j=1dAij𝔼[zizj]=j=1dAijδij=Aii

​ 由此我们证明了𝔼[z(Az)]=diag(A)

在神经网络中,相较于直接计算完整的Hessian矩阵以及其OBD方式的对角近似,计算其矩阵向量积(HVP,Hessian Vector product)是并不困难的,它只需要一次反向传播:

Hz=(gz)θ=gθz+gzθ=gθz

由此,我们可以通过多次在满足Rademacher分布的向量取样计算z(Az)的期望就能得到diag(A)无偏估计,在实际的应用中,只进行一次取样就能得到较为不错的结果。

外积近似

在神经网络应用于回归问题时,通常采用下面形式的平方和误差

E=12n=1N(yntn)2

​ 那么Hessian矩阵可以写成如下形式:

H=E=(n=1N(yntn)yn)=n=1Nyn(yn)+n=1N(ynt)yn

​ 在网络已经训练好的情况下,输出yntn接近,因此第二项可以忽略,由此我们得到了Hessian矩阵的外积近似

Hn=1Nbnbn

​ 其中,bn=yn=an(输出单元的激活函数就是恒等函数)。这种方法中的Hessian矩阵可以跟随反向传播算法在O(W)个步骤内高效地求出误差函数地一阶导数。再通过简单地乘法就可以在O(W2)步骤内求出矩阵元素。

Hessian矩阵逆的计算

使用外积近似,我们可以提出一个计算Hessian矩阵的逆的高效办法,首先我们有:

𝐇N=i=1nbnbn

其中,bn=wan时数据点n产生的输出单元激活对梯度的贡献。我们现在推导一个建立Hessian矩阵的顺序步骤,每次处理一个数据点。假设我们已经使用了前L个数据点得到了Hessian矩阵的逆。通过将第L+1个数据点的贡献单独写出来,我们有:

𝐇L+1=𝐇L+bL+1bL+1

为了计算Hessian矩阵的逆,我们考虑下面这个矩阵的恒等式:

(M+vv)1=M1(M1v)(vM1)1+vM1v

若我们令HL=M,且bL+1=v,我们有:

HL+11=HL1HL1bL+1bL+1HL11+bL+1HL1bL+1

通过这种方式,数据点可以依次使用,直到L+1=N,整个数据集处理完毕。于是,这个结果表示一个计算Hessian矩阵的逆的算法。这个算法只需对数据集扫描一次。最开始的矩阵H0被选为αI,其中α是一个较小的量,从而算法实际找的是𝐇+αI的逆。结果对于α的精确值不敏感。

Hessian矩阵的逆计算的K-FAC分解方法

​ 在特定场景下,Hessian矩阵的逆有一个利用Kronecker分解的近似方法,它利用了在loss为log-like形式下𝔼[H]=F的特点,其中F为Fisher信息矩阵:

F=𝔼[logp(x)logp(x)]

其中上式中做了如下形式的简写:

logp(x)=θlog(p|θ)𝔼[p(x)]=Ex~p(x|θ)

首先我们有:

𝔼[logp(x)]=(logp(x))p(x)dx=logp(x)p(x)p(x)dx=logp(x)dx=1=0

那么当loss为log-like形式下时:

L=logp(x)

我们求它的Hessian:

2L=2logp(x)=p(x)p(x)=p(x)p(x)P(x)p(x)P(x)2p(x)=logp(x)logp(x)2P(x)p(x)

对该式求期望,我们有:

𝔼[H]=𝔼[logp(x)logp(x)]𝔼[2p(x)p(x)]=F2p(x)p(x)p(x)dx=F2p(x)dx=F2p(x)=F

由于θ本质为所有层𝐖的拼接,我们有:

dθ=θL(x)F=𝔼[L(x)L(x)]=𝔼[dθdθ]θ=[vec(W0),vec(W1),,vec(Wn)]

带入展开得到:

Fij=𝔼[vec(dWi)vec(dWj)]

ai,gi分别为第i层的前向输入和反向传播梯度,由反向传播算法有:

dWi=giai

带入有:

vec(dWi)=vec(giai)=aigi

这里可能不太直观,这是因为平常对kronecker积()的接触较少,具体而言,我们这样定义Kronecker积:

AB=[a11Ba1nBam1BamnB]

在这里,由于我们将dWiWi第i层参数的梯度进行了向量化,原本的Wi 是如下形式:

Wi=giai=[g1a1g1angma1gman]

按列堆叠(vec)就有:

vec(dWi)=[g1a1g1angman]=aigi

那么我们可以将Fisher矩阵写作如下形式:

Fij=𝔼[vec(dWi)vec(dWj)]=𝔼[(aigi)(ajgj)]=𝔼[(aiaj)(gigj)]𝔼[aiaj]𝔼[gigj]

这个近似相当于我们忽略了Cov(aiaj,gigj),这实际上是合理的,尤其是在网络较深的情况下。这里给出一个链接提供一个比较详细的说明。

虽然直接求Fij的复杂度仍然不变,但是利用Kronecker积的性质,我们在求逆时可以得到较大的性能提升:

(AB)1=A1G1

并且在实际运算中,并不是整个网络的Fisher进行计算,而是按层做块对角运算:

𝐅=blockdiag(F1,F2,),FlAlGl

这样每层求逆的时候只需要求两个小矩阵的逆。

一些统计量的计算

​ 除了直接对Hessian矩阵的近似计算,我们有时候仅仅需要对Hessian矩阵的统计量进行计算,例如矩阵的迹,最大特征值等。

​ 在计算这些统计量时,我们需要一个重要的算子:Hv,即Hessian矩阵与任意向量的乘积,这个乘积我们在上面已经证明了通过一次简单的反向传播算法可以得到。(PyHessian是一个Python库,它实现了这个算子的高效计算)

迭代:

yk+1=Axkxk+1=yk+1||yk+1||μk+1=xk+1Axk+1xk+1xk+1

具体的收敛性证明见教材

简单来说,方法一样,从Rademacher分布中取一随机向量v,然后有恒等式:

Tr(H)=Tr(HI)=Tr(H𝔼[vv])=𝔼[Tr(Hvv)]=𝔼[vHv]

Hessian矩阵在神经网络下的特殊结构

本小节内容主要基于Towards Quantifying the Hessian Structure of Neural Networks,B站上有作者的讲解视频[FAI] 港中深 张雨舜 | 浅谈神经网络Hessian矩阵的特殊结构

​ 这篇文章主要说明了在神经网络中,Hessian矩阵往往具有近块对角结构,这表明曲率信息(二阶信息)主要在层内耦合,层与层的二阶交互很弱。这间接的说明了Layer-wise的量化/剪枝的合理性。

​ 此外,文章还进一步从理论和实验上阐释了造成这个现象的原因:

  1. 静态力量:即使在随机初始化阶段,即在训练开始之前,神经网络的Hessian矩阵就已经呈现出近块对角结构。这种结构的形成与网络的架构设计有关,因此被称为“静态力量”。具体来说,对于线性模型和单隐藏层网络,无论是使用均方误差(MSE)损失还是交叉熵(CE)损失,Hessian矩阵的对角块和非对角块在随机初始化时就已经表现出明显的差异。
  2. 动态力量:在训练过程中,Hessian矩阵的结构会进一步发生变化。特别是在使用CE损失时,训练过程会逐渐消除初始时存在于跨层Hessian分量(Hwv)中的“块循环”(block-circulant)模式,而对角块和非对角块的近块对角结构则保持稳定。这种由训练过程引起的结构变化被称为“动态力量”。

另外一点,文章证明了类别数C是影响Hessian矩阵结构的主要因素。(在实验中,隐藏层Hessian的非对角块与对角块的比值以1C的速度衰减,输出层Hessian的衰减速率为1C).这个结果对于大模型而言是友好的,因为在神经网络中C的大小往往是1e3~1e4级别的,这说明大模型的Hessian是具有强对角块的结构的!这一点实际上在传统的方法上人们意识or无意识的用到了(对角近似),但是更丰富层面以及基于这个发现的在计算的可行性和性能的Trade-off做的工作是比较少的,在优化器那边做的比较多。

此外还有不少文献揭露了Hessian矩阵具有低秩特征谱,即只有少数的特征值显著大,其余大多接近0.