Low-rank Compression of Neural Nets:Learning the Rank of Each Layer 阅读笔记

论文地址: paper
代码地址: code

1. 论文概述

可以通过使用低秩矩阵近似逼近每层权重的方法实现神经网络的压缩，但难点在于每层的最佳秩都是一个超参搜索的问题。针对上述问题，本片文章基于秩和矩阵元素提出了一种混合离散-连续优化函数。

本文提出了一种近似解决这个问题的算法，首先针对该问题的描述建立在减小分类网络的误差和基于秩的模型选择损失，利用秩约束网络的每层卷积。然后这个问题可以通过秩和权重进行优化，交错使用SGD步骤来训练未压缩网络和确定当前最优秩和权重矩阵。

2. 问题描述

2.1 变量定义

$K$ 表示网络的层数，$k$ 表示网络层中第$k$层，$z_k=\sigma (W_k{z}_{k-1})$, 其中 $\sigma (\cdot )$表示激活函数, $z_k$是$k$层的输入($x=z_0和y=z_K$表示网络的输入和输出)，$L(W)$表示损失函数，其中 W = { W 1 , W 2 , ⋯   , W K } W=\left\{W_1, W_2, \cdots, W_K\right\} W={W1​,W2​,⋯,WK​}是权重矩阵，$W_k$的尺寸是$a_k\times b_k$, 其中$b_k=a_{k-1}$。因此，本文章提出的目标函数是优化损失函数和代价函数$C$的和，其中限制条件是每层的秩小于最大的秩 $R_k$, $R_k\le min(a_k,b_k)$.
$$\underset{W}{\min }L(W)+\lambda C(W)$$
$$s.t.rank(W_k)\le R_k,k=1,2,\dots ,K.$$
其中上述约束条件等价于$W_k=U_k{V}_k^T$, $U_k$的尺寸是$a_k\times b_k$, $V_k$的尺寸是$b_k\times r_k$, r k ∈ { 0 , 1 , … , R k } r_k \in \left\{0,1,\dots,R_k\right\} rk​∈{0,1,…,Rk​}, 但是 $r_k$是一个待优化的未知参数，因此上述目标函数可以转为联合优化 { W k , U k , V k , r k } k = 1 K \left\{W_k, U_k, V_k, r_k\right\}_{k=1}^{K} {Wk​,Uk​,Vk​,rk​}k=1K​，$\lambda$是分类损失函数的权重系数，代价函数$C(W)$的定义如下：
$$C(W)=C(r_1,\dots ,r_K)=C_1(r_1)+\cdots +C_K(r_K)$$
特殊案例是$C(W)={\alpha }_1{r}_1+\cdots +{\alpha }_K{r}_K$, 其中 ${\alpha }_1,\cdots ,{\alpha }_K\ge 0$的常数，从优化的角度看， $C$仅取决于秩，从建模的角度看通过选择适当的系数${\alpha }_k$, $C$可以表示模型压缩中的代价函数。

2.2 优化策略

定义辅助变量$\Theta =({\Theta }_1,\dots ,{\Theta }_K)$ 和限制条件 $\Theta =W$, 上述目标函数可以改写为：
$$\underset{W,\Theta ,r}{\min }L(W)+\lambda C(r)$$
$$s.t.W_k={\Theta }_k,rank({\Theta }_k)=r_k\le R_k,k=1,\dots ,K,r=(r_1,\dots ,r_k)$$

使用学习压缩策略(learning-compression, LC)对上述目标函数进行优化, 应用 quadratic-penalty 方法之后优化目标变为：
$$Q(W,\Theta ,r;\mu )=L(W)+\lambda C(r)+\frac{\mu }{2}\sum _{k=1}^K||W_k-{\Theta }_k|{|}^2$$
$$s.t.rank({\Theta }_k)=r_k\le R_k,k=1,\dots ,K$$

上述目标中的待优化的参数$W$和$(\Theta ,r)$。优化$W$的过程称为$L$阶段，其目标函数是${\min }_{W}L(W)+\frac{\mu }{2}{\sum }_K||W_k-{\Theta }_k|{|}^2$。优化 $\Theta$和$r$的过程称为$C$阶段，目标函数如下：
$$\underset{{\Theta }_k,r_k}{\min }\lambda C_k(r_k)+\frac{\mu }{2}||W_k-{\Theta }_k|{|}^2$$
$$s.t.rank({\Theta }_k)=r_k\le R_k$$

假设 $a_k\ge b_k$ 和 $W_k=U_k{S}_k{V}_k^T$, 其中 $U_k$的尺寸$a_k\times b_k$和$V_k$的尺寸$b_k\times b_k$, $S_k=diag(s_1,\dots ,s_{b_k}),s_1\ge \cdots \ge s_{b_k}\ge 0$。则目标函数等价于：
min ⁡ r λ C k ( r ) + μ 2 ∑ i = r + 1 R k s k i 2 , s . t . r k ∈ { 0 , 1 , … , R k } \min_r \lambda C_k(r)+\frac{\mu}{2}\sum_{i=r+1}^{R_k}s^2_{ki}, s.t. r_k \in \left\{0,1,\dots,R_k\right\} rmin​λCk​(r)+2μ​i=r+1∑Rk​​ski2​,s.t.rk​∈{0,1,…,Rk​}

具体的算法流程图如下：

3. 文章总结

本文章的主要出发点是从分类损失函数和权重的代价函数两个角度出发构建目标函数，同时使用秩对权重的代价函数进行限制。