大模型压缩：基于贝叶斯优化的自适应低秩分解

1.方法

1.1 基于特征的高维空间低秩分解

PCA已经是老朋友了，每次一说主成分都会出现PCA。这篇文章¹利用预训练数据的子集作为校准数据集 D c a l = { x i } i = 1 n \mathcal{D}_{cal}=\{x_{i}\}_{i=1}^{n} Dcal​={xi​}i=1n​，首先用校准数据集的样本协方差矩阵（SCM）估计整个特征空间分布的Y的协方差矩阵
$$Co{v}_S(\mathbf{Y})=\frac{1}{n-1}\sum _{i=1}^n({\mathbf{y}}_i-\bar{\mathbf{y}}{)}^T({\mathbf{y}}_i-\bar{\mathbf{y}})\text{\hspace{1em}(1)}$$
式中${\mathbf{y}}_i$表示${\mathbf{x}}_i$的特征，$\bar{\mathbf{y}}$是校准数据集的特征值平均值。但文章指出，计算高维的协方差矩阵并不简单，他们提出了合并协方差矩阵（PCM），把校准数据集分成$m$组，对每一组分别计算协方差矩阵最后求平均得PCM
$$Co{v}_P(\mathbf{Y})=\frac{1}{m}\sum _{k=1}^{m}Co{v}_S({\mathbf{Y}}_k)\text{\hspace{1em}(2)}$$

1.2 基于贝叶斯优化得低秩分配

对于每一层，甚至每一层的不同矩阵对低秩分解得敏感度不同，对于一个模型$f(\cdot ;\theta )$，引入一组压缩率 λ = { λ i } i = 1 k \lambda=\{\lambda_{i}\}_{i=1}^{k} λ={λi​}i=1k​，引入一个任务模糊数据集D来评价压缩大模型$f(\cdot ;\mathbf{\theta },\lambda )$的性能，例如预训练数据集的子集。因此目标函数表示为
$$\begin{array}{rl}\underset{\lambda \in \mathcal{V}}{\min }H(\mathbf{\lambda })& ={\mathbb{E}}_{(x,y)\sim \mathcal{D}}h(f(x;\mathbf{\theta },\mathbf{\lambda }),y)\\ & s.t.\Sigma \mathbf{\lambda }\le \rho \end{array}\text{\hspace{1em}(3)}$$
式中，$\rho$是模型的整体压缩比，$h(\cdot ,\cdot )$是评价指标，但事实上，评价指标和低秩分配都是非常耗时耗算力的，所以这篇论文提出样本高效贝叶斯优化（BO）来优化公式3。这里引入一个替代模型（例如高斯模型$\mathcal{N}(\mu (\cdot ),{\sigma }^2(\cdot ))$），BO通过替代模型来估计目标函数$H(\mathbf{\lambda })$，并且基于每一步的结果，更新后面一步的目标函数$H(\mathbf{\lambda })$。比如给出前t-1步 { λ 1 , ⋯   , λ t − 1 } \{\boldsymbol{\lambda}_{1},\cdots,\boldsymbol{\lambda}_{t-1}\} {λ1​,⋯,λt−1​}的目标函数值分别为$H_{t-1}=[H({\mathbf{\lambda }}_1),\cdots ,H({\mathbf{\lambda }}_{t-1})]$，替代模型更新为$$\begin{gathered} \mu (\mathbf{\lambda })=\mathbf{k}(\mathbf{K}+{\eta }^2\mathbf{I}{)}^{-1}{H}_{t-1} \\ {\sigma }^2(\mathbf{\lambda })=k(\mathbf{\lambda },\mathbf{\lambda })-{\mathbf{k}}^T(\mathbf{K}+{\eta }^2\mathbf{I}{)}^{-1}\mathbf{k}\text{\hspace{1em}(4)} \end{gathered}$$
式中$k(\cdot ,\cdot )$是一个核函数，$(\mathbf{k}=k(\mathbf{\lambda },{\mathbf{\lambda }}_i){)}_{i\in [t-1]}$，$K=(k({\mathbf{\lambda }}_i,{\mathbf{\lambda }}_j){)}_{i,j\in [t-1]}$，${\eta }^{2}I$是用来模拟噪声的白核函数，得到后验估计$H(\mathbf{\lambda })$（例如$H(\mathbf{\lambda })\sim \mathcal{N}(\mu (\mathbf{\lambda }),{\sigma }^2(\mathbf{\lambda }))$）之后，BO通过采集函数确定下一次的比率分布，对于采集函数，可以用EI
α ( λ ) = E H ( λ ) [ max ⁡ { 0 , H ′ − H ( λ ) } ] λ t = a r g m a x λ α ( λ ) , (5) \begin{aligned}\alpha(\boldsymbol{\lambda})&=\mathbb{E}_{H(\boldsymbol{\lambda})}\left[\max\left\{0,H'-H(\boldsymbol{\lambda})\right\}\right]\\\boldsymbol{\lambda}_{t}&=\mathop{\mathrm{argmax}}_{\boldsymbol{\lambda}}\alpha(\boldsymbol{\lambda}),\end{aligned}\tag{5} α(λ)λt​​=EH(λ)​[max{0,H′−H(λ)}]=argmaxλ​α(λ),​(5)
式中，$H^{\prime }={\min }_{i\in [t-1]}H({\mathbf{\lambda }}_i)$是指迄今为止观察到的最小值，然后BO选择了最好的EI的方向去搜索。在得到最优比${\lambda }^{\ast }$之后，可以确定分配$r_i=(1-{\lambda }_i)d_1{d}_2/(d_1+d_2)$。

1.3 后训练

为了不使模型参数量反弹，文章使用压缩模型的子空间对模型微调。
$$Y=(BA+{\Lambda }_b{B}_{r^{\prime }}{\Lambda }_d{A}_{r^{\prime }})X\text{\hspace{1em}(6)}$$
式中，$B_{r^{\prime }}\in {\mathbb{R}}^{d_2\times r^{\prime }}$，$A_{r^{\prime }}\in {\mathbb{R}}^{r^{\prime }\times d_1}$是修正后的$B$和$A$矩阵，${\mathbf{\Lambda }}_b$和${\mathbf{\Lambda }}_d$是对角阵。

---

1. 基于贝叶斯优化的自适应低秩分解 ↩︎