论文笔记 SliceGPT: Compress Large Language Models By Deleting Rows And Columns

欲买桂花同载酒，终不似，少年游。

数学知识

秩： 矩阵中最大线性无关的行/列向量数。行秩与列秩相等。
线性无关：对于N个向量而言，如果任取一个向量$\text{v}$，不能被剩下的N-1个向量通过线性组合的方式表示，则称这N个向量为线性无关。

SliceGPT: Compress Large Language Models By Deleting Rows And Columns

主要剪枝的效果：

1. 将权重矩阵的尺寸缩小，变成更小的矩阵。具体而言是乘以一个Deleting Matrix：$\text{D}$。这个矩阵实际上是通过构造一个正交矩阵，再做PCA删除一些行/列得到。
2. 减少embedding dimension。与权重矩阵的缩小对应。
   可以参考Figure 1中右图

0. Introduction

目前许多权重剪枝的方法都需要RFT（recovery fine-tuning），耗时并且可拓展性差。

• SliceGPT无需RFT也能有良好的效果。

本文的三大Contributions

1. 引入了计算不变性。对于transformer中的权重矩阵做正交变换。简而言之就是乘上一个正交矩阵和正交矩阵的转置，计算结果不变。
2. 使用signal matrix计算正交阵，利用PCA，在其主成分方向做投影后，与权重矩阵相乘，移去部分columns和rows以达到权重剪枝的目的。
3. 在多个模型，不同任务上做实验并证明效果良好。

1. Related Work

常见的稀疏化方法：

1. magnitude-based：移除绝对值较小的权重，但性能损失较大。
2. OBS（optimial brain surgeon）：使用Hessian矩阵来更新weight，移除对loss函数影响最小的weight。但是计算Hessian的逆过于复杂，尤其是大模型。
3. 针对OBS的问题：
   1. 近似Hessian，如WoodFisher
   2. 逐层使用OBS，比如OBC（optimal brain compression）。
4. GPTQ：量化权重矩阵。
5. sparseGPT：半结构化/非结构化剪枝
6. 仅使用Hessian的对角项（对角近似）

2. Transformer模型架构回顾。

相关信息可以参阅CSDN或知乎关于Transformer原文的详解，也可以参考我的这篇文章（内容比较简略）：https://blog.csdn.net/qq_43349296/article/details/144001387

一个典型的Transformer中Encoder的结构Figure 2所示，作者把这样一个结构称为Transformer Block（不同的Transformer实现可能有所不同）：

简要回顾一下Transformer的主要组件，以及约定符号表示。
一个标准的Transformer Block包含Attention Block，FFN Block，LayerNorm层。Transformer在起始阶段有embedding层将sequence转为嵌入向量，末尾有一个输出层,文中称为Language Modelling Head，将embedding还原为sequence。

• Embeddings层：将sequence：$S$，转为embedding：$\text{X}\in R^{N\times D}$。其中$N$为序列长度，$D$为嵌入向量的长度，同时一般也是模型当中统一的hidden dimension。$W_{\text{embd}}\in R^{h\times D}$为该层的权重矩阵。$h$为one-hot编码长度。

• LN层（LayerNorm）：标准的Transformer使用LN，而作者使用RMSNorm原因是其具有计算不变性（Computational Invariance）。RMSNorm相比于LN而言计算更简单，每个元素只除以RMS（均方根，并且这里求均方根时，统计的是$\text{X}$中的所有元素: 共$N\times D$个，这也是为什么后面要乘以一个$\sqrt{D}$的原因。LN是对每个样本做归一化，即按行归一化）即可，LN与RMSNorm之间的关系如下：
  $$\text{LayerNorm}(\mathbf{X})=\text{RMSNorm}(\mathbf{XM})\text{diag}(\alpha )\sqrt{D}+1_N{\beta }^{\top }\text{\hspace{1em}(1)}$$

其中：
$$\text{M}=I-\frac{1}{D}1\cdot 1^{T}s.t.I\in {\mathbb{R}}^{D\times D},1\in {\mathbb{R}}^{D\times 1}$$

输入乘上$\text{M}$相当于逐行减去mean。$diag(\alpha )$为缩放系数，$\beta$为偏置项。

• Attention Blocks：使用多头注意力，${\mathbf{W}}_k,{\mathbf{W}}_q,{\mathbf{W}}_v,{\mathbf{W}}_o,$分别对应K, Q, V, Output的权重矩阵。（output是一个Linear，把各个head拼接后的embedding再映射回去）。Attention Block用以下公式表示：
  $$\sigma ({\mathbf{W}}_{\mathrm{in}}+{\mathbf{b}}_{\mathrm{in}}){\mathbf{W}}_{\mathrm{out}}+{\mathbf{b}}_{\mathrm{out}}\text{\hspace{1em}(2)}$$
  作者把${\mathbf{W}}_k,{\mathbf{W}}_q,{\mathbf{W}}_v$统称为$W_{in}$，因为这几个矩阵对attention block的输入做线性变换，把${\mathbf{W}}_o$称为${\mathbf{W}}_{out}$，因为是attention中的输出层（多头注意力中把concatenated embedding映射回原维度）。

• FFN Blocks: $\sigma ({\text{XW}}_{\text{in}}){\text{W}}_{\text{out}}$。即MLP，简而言之就是先后做两次线性变换，先升维，再还原维度。

• LM Head： ${\text{XW}}_{\text{head}}+{\text{b}}_{\text{head}}$。其中$\text{X}$为最后一个FFN Block的输出。LM Head输出即为最终的预测word。

Transformer整体的前向传播流程如Algorithm1所示：

3.SliceGPT

Key Idea：Computational Invariance。即：对线性层（使用nn.Linear的层如Attention，FFN）施加正交变换，计算结果不变。

3.1 Transformer中的computational invariance的说明

• 正交矩阵的保范性：假设$\text{Q}$为正交矩阵，则$\text{Q}{\text{Q}}^T=\text{I}$，对于向量$\text{x}$，$||\text{Q}\text{x}||=\sqrt{{\text{x}}^T{\text{Q}}^T\text{Q}\text{x}}=\sqrt{{\text{x}}^T\text{x}}=||\text{x}||$。即向量乘以正交阵不改变其范数。这里列出的是L2范数。

作者指出RMSNorm具有计算不变性，如eq2所示，作者在Appendix A.1给出了证明：
（列出公式2）

3.1（续）定理一以及证明：

定理一：作者指出，给Transformer当中的权重矩阵施加正交变换，能够保证其计算不变性：
$$\begin{array}{crcr}{\tilde{\mathbf{W}}}_{embd}={\mathbf{W}}_{embd}\mathbf{Q},& \text{(3)}& {\tilde{\mathbf{b}}}_{out}^{\ell }={\mathbf{b}}_{out}^{\ell }\mathbf{Q},& \text{(6)}\\ {\tilde{\mathbf{W}}}_{in}^{\ell }={\mathbf{Q}}^{\top }{\mathbf{W}}_{in}^{\ell },& \text{(4)}& {\tilde{\mathbf{W}}}_{head}={\mathbf{Q}}^{\top }{\mathbf{W}}_{head}.& \text{(7)}\\ {\tilde{\mathbf{W}}}_{out}^{\ell }={\mathbf{W}}_{out}^{\ell }\mathbf{Q},& \text{(5)}& & \end{array}$$
加波浪线的为变换后（microsoft实现代码中称为rotate 即旋转）

• 注：原文中eq.6为$\tilde{\boldsymbol{b}}{out}^{{\ell}=\mathbf{Q}}\top\boldsymbol{b}{out}^{\ell} $，好像有问题，正在向作者咨询。
• 注：${\tilde{\mathbf{b}}}_{in}^{\ell }={\mathbf{b}}_{in}^{\ell },{\tilde{\mathbf{b}}}_{head}={\mathbf{b}}_{head}.$

在这里我们也可以简单证明一下。我们参考Algorithm 1中1-7行，对前向传播的各个步骤给出对应的公式：
$$\begin{array}{rl}& 1:\mathbf{X}\leftarrow S{\mathbf{W}}_{\mathrm{embd}}\\ & 2:\mathbf{X}\leftarrow {\mathrm{RMSNorm}}_0(\mathbf{X})\\ & 3:\mathbf{for}\ell =1\dots L\mathbf{do}\\ & 4:\mathbf{Z}\leftarrow {\sigma }_{\ell }({\mathbf{XW}}_{\mathrm{in}}^{\ell }+{\mathbf{b}}_{\mathrm{in}}^{\ell }){\mathbf{W}}_{\mathrm{out}}^{\ell }+{\mathbf{b}}_{\mathrm{out}}^{\ell }\\ & 5:\mathbf{X}\leftarrow {\mathrm{RMSNorm}}_{\ell }(\mathbf{X}+\mathbf{Z})\\ & 6:\text{ end for}\\ & 7:{\text{ return XW}}_{\mathrm{head}}+{\mathbf{b}}_{\mathrm{head}}\end{array}$$
其中$S\in R^{N\times h}，{\mathbf{W}}_{\mathrm{embd}}\in R^{h\times D}，{\mathbf{W}}_{\mathrm{head}}\in R^{D\times h}$。为了简化，统一认为${\mathbf{W}}_{\mathrm{in}},{\mathbf{W}}_{\mathrm{out}}\in R^{D\times D}$。其中$N$为序列长度，$h$为one-hot编码的长度，$D$为hidden dimension（或者叫embedding dimension）。

施加正交矩阵$\mathbf{Q}$后的各步骤公式如下，我们将$\mathbf{X},\tilde{\mathbf{X}}$分别表示为施加正交变换前，正交变换后block的输入/输出：

• $\text{line1}:S{\tilde{\mathbf{W}}}_{\mathrm{embd}}=S{\mathbf{W}}_{\mathrm{embd}}\mathbf{Q}=\mathbf{XQ}\to \tilde{\mathbf{X}}$

• $\text{line2}:\text{RMSNorm}(\tilde{\mathbf{X}})=\text{RMSNorm}(\mathbf{XQ})=\text{RMSNorm}(\mathbf{X})\mathbf{Q}\to \tilde{\mathbf{X}}$

• $\text{line4}:{\sigma }_{\ell }({\tilde{\mathbf{X}}\tilde{\mathbf{W}}}_{\mathrm{in}}^{\ell }+{\tilde{\mathbf{b}}}_{\mathrm{in}}^{\ell }){\tilde{\mathbf{W}}}_{\mathrm{out}}^{\ell }+{\tilde{\mathbf{b}}}_{\mathrm{out}}^{\ell }={\sigma }_{\ell }({\mathbf{XQ}{\mathbf{Q}}^{\top }\mathbf{W}}_{\mathrm{in}}^{\ell }+{\mathbf{b}}_{\mathrm{in}}^{\ell }){\mathbf{W}}_{\mathrm{out}}^{\ell }\mathbf{Q}+{\mathbf{b}}_{\mathrm{out}}^{\ell }\mathbf{Q}=({\sigma }_{\ell }({\mathbf{XW}}_{\mathrm{in}}^{\ell }+{\mathbf{b}}_{\mathrm{in}}^{\ell }){\mathbf{W}}_{\mathrm{out}}^{\ell }+{\mathbf{b}}_{\mathrm{out}}^{\ell })\mathbf{Q}=\mathbf{ZQ}\to \tilde{\mathbf{Z}}$

• $\text{line5}:\text{RMSNorm}(\tilde{\mathbf{X}}+\tilde{\mathbf{Z}})=\text{RMSNorm}(\mathbf{XQ}+\mathbf{ZQ})=\text{RMSNorm}(\mathbf{X}+\mathbf{Z})\mathbf{Q}=\mathbf{XQ}\to \tilde{\mathbf{X}}$

• $\text{line7}:\tilde{\mathbf{X}}{\tilde{\mathbf{W}}}_{\text{head}}+{\tilde{\mathbf{b}}}_{\text{head}}=\mathbf{XQ}{\mathbf{Q}}^{\top }{\mathbf{W}}_{\text{head}}+{\mathbf{b}}_{\text{head}}=\mathbf{X}{\mathbf{W}}_{\text{head}}+{\mathbf{b}}_{\text{head}}\to \tilde{\mathbf{X}}=\mathbf{X}$

发现 $\text{line7}$ 结果相等。综上，可以证明变换前后输出不变。

3.2 Transformer中LN向RMSNorm的转换：

根据eq.1我们可以知道LN与RMSNorm之间存在转换关系。其中最重要的两个就是mean-substraction: $\mathbf{M}$，以及系数: $\text{diag}(\alpha )$。作者指出，可以将LayerNorm中的这两个步骤分别放在前一个Block与后一个Block当中，如Figure 3所示。可以对比一下Figure 2与Figure 3有哪些不同。

可以发现，${\mathbf{W}}_{\text{in}}$都是左乘 $\text{diag}(\alpha )$，而${\mathbf{W}}_{\text{out}}$均为右乘$\mathbf{M}$。除了考虑Figure 3中所包含的Attention Block以及FFN Block当中的${\mathbf{W}}_{\text{in}}$与${\mathbf{W}}_{\text{out}}$以外，考虑${\mathbf{W}}_{\text{embd}},{\mathbf{W}}_{\text{head}}$应该分别左乘$\text{diag}(\alpha )$、右乘$\mathbf{M}$。（这里其实很好理解，因为embedding层位于第一个LN层的前面，而LM Head层恰好在最后一个LN层的后面）

用矩阵运算求均值：乘以矩阵$M$即可。对最后一个维度求均值（对一行求均值）：
$$M=I-\frac{1}{D}1\cdot 1^{T}s.t.I\in {\mathbb{R}}^{D\times D},1\in {\mathbb{R}}^{D\times 1}$$
因此严格来说，将LayerNorm中的均值相减操作融合至前一个block后，似乎与原始的模型不太一致，因为矩阵乘法不遵循交换律。但代码实现中直接对权重矩阵做了mean-substraction操作。本人理解可能是作者为了简便，以及希望可以pre-compute ${\mathbf{W}}_{\text{in}}$ 的一种权宜之计。（这里加粗处存疑，如有问题请大佬指正）

3.3 Transformation Per Block

作者指出，对不同的Block，应该根据当前输入的signal matrix的不同，计算得到不同的正交阵$\mathbf{Q}$。但是Algorithm 1中 $\text{line 5}$会存在等式不相等的情况：
$$(\tilde{\mathbf{X}}+\tilde{\mathbf{Z}})=(\mathbf{X}{\mathbf{Q}}_{\ell -1}+\mathbf{Z}{\mathbf{Q}}_{\ell })\ne (\mathbf{X}+\mathbf{Z}){\mathbf{Q}}_{\ell }.\text{因为不同Block正交阵不相等}$$
本质原因是存在Residual Connection。故每一个残差连接对应的$\text{X}$应当右乘${\mathbf{Q}}_{\ell -1}{\mathbf{Q}}_{\ell }$，以保证$\text{line 5}$等式成立。

最终的经过变换后的Transformer Block示意图如Figure 4所示：

3.3(续)如何构造正交阵Q

作者提出根据每一层不同的signal matrix，分别构造不同的正交阵。公式如下：
$${\mathbf{C}}_{\ell }=\sum _i{\mathbf{X}}_{\ell ,i}^{\top }{\mathbf{X}}_{\ell ,i}\text{\hspace{1em}(8)}$$
其中${\mathbf{X}}_{\ell ,i}$表示第$\ell$个$\text{RMSNorm}$层对于第$i$个sequence的输出。${\mathbf{Q}}_{\ell }$即为${\mathbf{C}}_{\ell }$经过特征分解后，按特征值从大到小排列的所有特征向量所构成的矩阵。
注：${\mathbf{C}}_{\ell }$为对称矩阵，有什么意义？首先实对称矩阵的特征值肯定为实数。

3.4 Slicing

类似PCA当中的操作，选取$\mathbf{Q}$的特征值最大的$D_{\text{small}}$个特征向量，构造删除矩阵$\mathbf{D}\in R^{D\times D_{\text{small}}}$，将$\mathbf{X}$映射为一个低纬度的特征$\mathbf{Z}$，然后再经过正交阵的转置，又变换回$\tilde{\mathbf{X}}$，相当于reconstruction的过程。如以下公式所示：

$$\mathbf{Z}=\mathbf{XQD},\tilde{\mathbf{X}}=\mathbf{Z}{\mathbf{D}}^{\top }{\mathbf{Q}}^{\top }.\text{\hspace{1em}(9)}$$

具体的slice过程如下图所示（Figure 1的右图）

多头注意力机制实现方法：

1. 第一种实现：将输入向量降维成多个低维向量，比如8个头，embedding维度为512，那么就有24个Linear(512, 64)，其中8个作为$W_q$，8个作为$W_k$，8个作为$W_v$，这里Linear(512, 64)即是权重矩阵，也起到降维作用。然后8个降维后的向量各自做各自的attention，得到attention中每个head的输出（每个维度为64），再把这8个输出拼接起来，得到维度为512，然后再经过一个线性层Linear(512,512)，得到multi-head attention最终的输出。
2. 第二种实现：还是8个头的注意力机制，但是$W_q，W_k，W_v$都只有一个，为Linear(512,512)，然后将embedding reshape（使用view函数），把shape变换为[N, seq, head_num, head_dim]分别对应为样本数，序列中单词长度，head数量，每个head分得的维度数（比如8个头，则shape为[N, seq, 8, 64]），然后直接做点乘，最后再reshape将维度变换回去，再经过线性层Linear(512,512)，得到多头注意力机制的最终输出。

问题

1. 如何做Fusion？即如何将LN中的Linear operation融合至相邻线性层？
   Embedding层只做了mean substraction。其余层直接和LN层的参数做element-wise multiplication。相当于乘以了缩放系数，这里没问题。
2. 如何求解正交矩阵$\mathbf{Q}$？对signal matrix（指的是input / embedding）使用PCA，QR分解（后续需补充QR分解和特征分解的关系）。
3. 文中提到哪些部分不能pre-computed？指的是${\mathbf{Q}}_{\ell -1}{\mathbf{Q}}_{\ell }$，可以从代码实现中看到，该算法是一边forward，一边剪枝，也就是需要等当前的Block前向传播完毕后，拿到当前Block的输出（下一个Block的输入），才能开始计算下一个Block的$\mathbf{Q}$。比如当前是第$\ell -1$个Block，等这一个Block前向传播完毕后，才能开始算${\mathbf{Q}}_{\ell }$。
4. Norm层的可学习参数是指的哪些？指$\gamma ，\beta$即缩放系数和偏移量: https://www.cnblogs.com/tian777/p/17911800.html
5. Convolutional Layer是否也具有计算不变性？
6. LN层为什么不具备计算不变性？
7. 解释RMSNorm为什么具备计算不变性？

词语释义

cornerstone 基石
post-hoc = after this 事后的adj./事后adv.
complementary 补充的
undertaking 任务/项目
so long as 只要
whilst 与此同时

参考链接

• SliceGPT原文链接
• SliceGPT源码
• LLM大模型压缩——ICLR 2024 SliceGPT（原理详解）
• SliceGPT概述
• Phi-2 Transformer模型代码
  (作者的实验代码中给出了Phi-2对应的ModelAdapter以及LayerAdapter的实现)
• 机器之心: SliceGPT