ShortGPT：揭示大型语言模型的层冗余与高效压缩

ShortGPT：揭示大型语言模型的层冗余与高效压缩

作为一名研究Transformer架构的学者，你可能已经注意到大型语言模型（LLMs）在性能提升的同时，参数规模也急剧膨胀，动辄达到数十亿甚至万亿级别。这不仅对计算资源提出了巨大挑战，也推动了模型压缩技术的发展。近期，来自Baichuan Inc.和ISCAS的研究团队发表了一篇题为《ShortGPT: Layers in Large Language Models are More Redundant Than You Expect》的论文，提出了一个新颖的视角：LLMs的层级冗余远超预期，并基于此开发了一种简单而高效的层剪枝方法——ShortGPT。本文将为熟悉Transformer的研究者介绍这篇论文的核心贡献，并深入解释其关键数学公式。

论文核心贡献

1. 揭示层级冗余的现象
   论文通过实验和理论分析，揭示了LLMs（尤其是基于预归一化pre-norm的Transformer架构）在层级上存在显著冗余。研究发现，某些层的输入与输出高度相似，对模型整体功能贡献较小。特别是在预归一化配置下，深层变换的作用被进一步弱化。这种冗余不仅限于Transformer架构，还在非Transformer模型（如Mamba和RWKV）中得到了验证，表明层级冗余可能是现代LLMs的普遍特性。

2. 提出Block Influence（BI）指标
   为了量化每层的贡献，论文引入了一个新颖的指标——Block Influence（BI）。BI通过计算层输入与输出的余弦相似度，衡量层对隐藏状态的变换程度。BI得分低的层被认为是冗余层，可以被移除。实验表明，BI指标与层的重要性高度相关，优于其他指标（如顺序、反序或相对幅度）。

3. 开发ShortGPT剪枝方法
   基于BI指标，论文提出了一种简单直接的剪枝方法——ShortGPT，通过移除BI得分较低的层来压缩模型。实验表明，ShortGPT能够在减少约25%参数的同时，保留约90%的性能。例如，在LLaMA 2-13B模型上，移除10层（占总层数的25%）后，MMLU基准性能仅从55.0降至52.2，优于其他复杂剪枝方法（如LLM-Pruner和SliceGPT）。

4. 与量化方法的正交性
   ShortGPT不仅在性能上表现出色，还与量化方法（如GPTQ）正交，可以结合使用进一步降低计算开销。实验显示，在量化后的LLaMA 2-7B模型上应用ShortGPT，模型性能和吞吐量均得到显著优化。

5. 后训练恢复性能
   为缓解层移除带来的性能损失，论文探索了后训练策略，包括用轻量级多层感知机（MLP）替换被移除的层并进行重新训练。这种方法有效恢复了部分性能，进一步提升了ShortGPT的实用性。

数学公式解析

论文的核心理论分析围绕预归一化（pre-norm）如何导致层输入与输出高相似度展开，并通过数学推导解释了层冗余的成因。以下是关键公式的详细解析：

Block Influence（BI）公式

BI指标用于衡量第 ($i$) 层的重要性，定义如下：

$${\mathrm{BI}}_i=1-{\mathbb{E}}_{X,t}\frac{X_{i,t}^T{X}_{i+1,t}}{\Vert X_{i,t}{\Vert }_2\Vert X_{i+1,t}{\Vert }_2}$$

• 解释：

  • ($X_{i,t}$) 表示第 ($i$) 层在时间步 ($t$) 的隐藏状态（即层的输入）。
  • ($X_{i+1,t}$) 表示第 ($i$) 层输出（即第 ($i+1$) 层的输入）。
  • ($\frac{X_{i,t}^T{X}_{i+1,t}}{\Vert X_{i,t}{\Vert }_2\Vert X_{i+1,t}{\Vert }_2}$) 计算输入与输出的余弦相似度，值越接近1，表示两者越相似，层的变换作用越小。
  • (${\mathbb{E}}_{X,t}$) 表示对输入样本 ($X$) 和时间步 ($t$) 的期望，考虑了多个样本的平均表现。
  • ($1-\text{相似度}$) 将相似度转化为“变换程度”，因此BI值越低，说明层的作用越小，越可能是冗余层。

• 意义：
  BI通过量化隐藏状态的变换程度，提供了一个直观的层重要性评估标准。实验表明，BI得分与移除层后模型性能下降（PPL升高）的程度呈正相关，验证了其有效性。

预归一化导致高相似度的理论推导

论文在附录A中通过数学推导解释了预归一化（以RMSNorm为例）为何导致深层输入与输出高相似度。核心推导如下：

1. 隐藏状态的模增长
   根据Xiong等人（2020）的引理，在预归一化Transformer初始化时，第 ($L$) 层隐藏状态 ($x_L$) 的模满足：

   $$\left(1+\frac{L}{2}\right)d\le \mathbb{E}(\Vert x_L{\Vert }_2^2)\le \left(1+\frac{M}{2}\right)d$$

   其中 ($d$) 是隐藏状态维度，($M$) 是模型参数的某一常数。假设 ($x_L$) 的每个分量均值为0，可得：

   $$\Vert x_L\Vert =\Theta (\sqrt{L})$$

   这表明随着层数 ($L$) 增加，隐藏状态的模逐渐增大。

2. 层变换的相对贡献
   假设第 ($L+1$) 层的隐藏状态为：

   $$x_{L+1}=x_L+f_L(x_L,{\theta }_L)$$

   其中 ($f_L(x_L,{\theta }_L)$) 表示注意力或MLP操作，(${\theta }_L$) 为可学习参数。论文证明，对于注意力机制，($f_L$) 的模为常数级别：

   $$\Vert f_L(x_L,{\theta }_L)\Vert =O(1)$$

   以注意力为例：

   $$f_L(x_L,{\theta }_L)=\left(\text{softmax}(Q^{T}K)\frac{X_L}{\Vert X_L\Vert }\cdot v_{\text{rms}}\right)W_v{W}_q$$

   其模为：

   $$\Vert f_L\Vert =O(\Vert v_{\text{rms}}\Vert \Vert W_v\Vert \Vert W_q\Vert )=O(1)$$

   表明 ($f_L$) 的贡献不随层数 ($L$) 增长。

3. 余弦相似度计算
   计算第 ($L$) 层输入与输出的余弦相似度：

   $$\cos \text{similarity}(x_{L+1},x_L)=\frac{x_{L+1}\cdot x_L}{\Vert x_{L+1}\Vert \Vert x_L\Vert }=\frac{\Vert x_L{\Vert }^2+f_L(x_L,{\theta }_L)\cdot x_L}{\Vert x_{L+1}\Vert \Vert x_L\Vert }$$

   分母近似为：

   $$\Vert x_{L+1}\Vert =\Vert x_L+f_L\Vert \approx \Vert x_L\Vert (\text{因}\Vert f_L\Vert \ll \Vert x_L\Vert )$$

   因此：

   $$\cos \text{similarity}\ge \frac{\Vert x_L{\Vert }^2}{\Vert x_{L+1}\Vert \Vert x_L\Vert }-\frac{\Vert f_L\Vert \Vert x_L\Vert }{\Vert x_{L+1}\Vert \Vert x_L\Vert }=\frac{\Vert x_L\Vert }{\Vert x_{L+1}\Vert }-\frac{\Vert f_L\Vert }{\Vert x_{L+1}\Vert }$$

   代入 ($\Vert x_L\Vert =\Theta (\sqrt{L})$) 和 ($\Vert f_L\Vert =O(1)$)，得到：

   $$\cos \text{similarity}=\Theta \left(\sqrt{\frac{L}{L+1}}\right)-O\left(\sqrt{\frac{1}{L+1}}\right)$$

   当 ($L$) 很大时，($\sqrt{\frac{L}{L+1}}\to 1$)，第二项趋于0，说明深层的输入与输出余弦相似度接近1，变换作用微弱。

• 意义：
  该推导表明，预归一化使深层变换的作用被“稀释”，导致层输出与输入高度相似。这种高相似度是层冗余的理论基础，支持了移除深层的可行性。

实验亮点

• 跨模型验证：ShortGPT在LLaMA 2（7B和13B）、Baichuan2（7B和13B）以及非Transformer模型（Mamba-2.8B和RWKV-7B）上均表现出色，证明了方法的通用性。
• 性能与效率兼得：在LLaMA 2-7B-GPTQ模型上，移除27.1%的层后，MMLU性能仅下降约2%，而吞吐量提升19%。
• 后训练优化：通过替换为轻量级MLP并重新训练，ShortGPT在保持低参数量的同时，进一步缩小了性能差距。

局限性与未来方向

尽管ShortGPT表现优异，但论文指出其对多选任务的泛化能力影响较大，具体原因尚未完全阐明。此外，RWKV模型的冗余程度低于Mamba和Transformer，提示不同架构的冗余特性需要进一步探索。未来研究可以结合ShortGPT与其他压缩技术（如稀疏化）或优化训练流程，设计更高效的LLMs。

总结

《ShortGPT》通过揭示LLMs的层级冗余，提出了一种简单而强大的剪枝方法，为Transformer研究者提供了新的压缩视角。其BI指标和数学推导清晰地解释了冗余的来源，而实验结果则展示了方法的高效性和普适性。对于关注模型压缩和高效推理的研究者来说，这篇论文无疑值得深入阅读和实践。

BI解释

---

为什么 ($X_{i,t}$) 是矩阵，而不是向量？

在论文中，($X_{i,t}$) 被描述为“第 ($i$) 层在时间步 ($t$) 的隐藏状态（即层的输入）”，并且公式中提到 ($X_{i,t}^T{X}_{i+1,t}$)。你可能会疑惑：语言模型的输入通常是一个序列，每个时间步的隐藏状态不应该是一个向量吗？为什么这里看起来像是矩阵操作？

解答：

在实际的 Transformer 模型实现中，隐藏状态 ($X_{i,t}$) 并不是单个时间步的向量，而是整个序列在第 ($i$) 层的隐藏状态矩阵。让我逐步解释：

1. Transformer 的输入和隐藏状态表示：

   • 在 Transformer 模型中，输入是一个序列（例如一个句子），假设序列长度为 ($T$)（即有 ($T$) 个 token），每个 token 的嵌入维度为 ($d$)。因此，输入层的输出（经过嵌入和位置编码后）是一个矩阵 ($X\in {\mathbb{R}}^{T\times d}$)，其中每一行 ($X_t\in {\mathbb{R}}^d$) 是第 ($t$) 个 token 的嵌入向量。
   • 每一层 Transformer 的输入和输出也是一个矩阵，形状同样为 ($T\times d$)。第 ($I$) 层的输入是 ($X_i\in {\mathbb{R}}^{T\times d}$)，输出是 ($X_{i+1}\in {\mathbb{R}}^{T\times d}$)。这里的 ($X_{i,t}$) 表示矩阵 ($X_i$) 的第 ($t$) 行，也就是第 ($i$) 层、第 ($t$) 个 token 的隐藏状态向量，形状为 ($1\times d$)。

2. 为什么公式中 ($X_{i,t}$) 是向量，但计算时涉及矩阵？：

   • 公式中的 ($X_{i,t}$) 确实是第 ($t$) 个 token 的隐藏状态向量（维度为 ($d$)），但为了高效计算，实际操作中通常会一次性处理整个序列的隐藏状态矩阵 ($X_i$)。因此，论文中提到的“($X_{i,t}^T{X}_{i+1,t}$)”实际上是对每一行（每个 token）计算余弦相似度，然后通过期望 (${\mathbb{E}}_{X,t}$) 取平均。
   • 具体来说，($\frac{X_{i,t}^T{X}_{i+1,t}}{\Vert X_{i,t}{\Vert }_2\Vert X_{i+1,t}{\Vert }_2}$) 是第 ($t$) 个 token 在第 ($i$) 层输入和输出之间的余弦相似度。(${\mathbb{E}}_{X,t}$) 表示对所有样本 ($X$) 和所有时间步 ($t$)（即序列中的所有 token）取平均。

3. 矩阵化的计算：

   • 在实际实现中，为了高效计算 BI，我们通常不会单独对每个 ($X_{i,t}$) 和 ($X_{i+1,t}$) 计算余弦相似度，而是对整个矩阵 ($X_i$) 和 ($X_{i+1}$) 进行批量操作：
     • ($X_i^T{X}_{i+1}\in {\mathbb{R}}^{T\times T}$) 计算的是所有 token 对之间的点积。
     • 但我们只关心同一 token 的相似度（即 ($X_{i,t}$) 和 ($X_{i+1,t}$)），所以只取对角线的值。
     • 归一化项 ($\Vert X_{i,t}{\Vert }_2\Vert X_{i+1,t}{\Vert }_2$) 也是对每一行分别计算模长后相乘。
   • 最后，通过平均操作（(${\mathbb{E}}_{X,t}$)），得到整个序列的平均余弦相似度。

总结：

($X_{i,t}$) 在公式中表示第 ($t$) 个 token 的隐藏状态向量（维度为 ($d$)），但实际计算时，($X_i$) 是一个 ($T\times d$) 的矩阵，包含整个序列的隐藏状态。公式中的操作是对每个 token 的隐藏状态计算余弦相似度后取平均。

---

通过一个例子解释 BI 公式的计算

假设我们有一个小型 Transformer 模型，序列长度 ($T=3$)（3 个 token），隐藏状态维度 ($d=2$)，我们来计算第 ($i$) 层的 BI 得分。

假设数据：

• 第 ($i$) 层的输入矩阵 ($X_i$)（形状 ($T\times d=3\times 2$)）：
  $$X_i=\left[\begin{array}{cc}1& 0\\ 0& 1\\ 1& 1\end{array}\right]$$

  • 第 1 个 token 的隐藏状态 ($X_{i,1}=[1,0]$)
  • 第 2 个 token 的隐藏状态 ($X_{i,2}=[0,1]$)
  • 第 3 个 token 的隐藏状态 ($X_{i,3}=[1,1]$)

• 第 ($i$) 层的输出矩阵 ($X_{i+1}$)（形状 ($3\times 2$)）：
  $$X_{i+1}=\left[\begin{array}{cc}0.9& 0.1\\ 0.1& 0.9\\ 0.95& 0.95\end{array}\right]$$

  • 第 1 个 token 的隐藏状态 ($X_{i+1,1}=[0.9,0.1]$)
  • 第 2 个 token 的隐藏状态 ($X_{i+1,2}=[0.1,0.9]$)
  • 第 3 个 token 的隐藏状态 ($X_{i+1,3}=[0.95,0.95]$)

步骤 1：计算每个 token 的余弦相似度

余弦相似度的公式为：
$$\text{CosSim}(X_{i,t},X_{i+1,t})=\frac{X_{i,t}^T{X}_{i+1,t}}{\Vert X_{i,t}{\Vert }_2\Vert X_{i+1,t}{\Vert }_2}$$

• 对于第 1 个 token (($t=1$))：

  • ($X_{i,1}=[1,0]$), ($X_{i+1,1}=[0.9,0.1]$)
  • 点积：($X_{i,1}^T{X}_{i+1,1}=(1\cdot 0.9)+(0\cdot 0.1)=0.9$)
  • 模长：($\Vert X_{i,1}{\Vert }_2=\sqrt{1^2+0^2}=1$), ($\Vert X_{i+1,1}{\Vert }_2=\sqrt{0.9^2+0.1^2}=\sqrt{0.81+0.01}=\sqrt{0.82}\approx 0.905$)
  • 余弦相似度：($\frac{0.9}{1\cdot 0.905}\approx 0.995$)

• 对于第 2 个 token (($t=2$))：

  • ($X_{i,2}=[0,1]$), ($X_{i+1,2}=[0.1,0.9]$)
  • 点积：($X_{i,2}^T{X}_{i+1,2}=(0\cdot 0.1)+(1\cdot 0.9)=0.9$)
  • 模长：($\Vert X_{i,2}{\Vert }_2=\sqrt{0^2+1^2}=1$), ($\Vert X_{i+1,2}{\Vert }_2=\sqrt{0.1^2+0.9^2}=\sqrt{0.01+0.81}=\sqrt{0.82}\approx 0.905$)
  • 余弦相似度：($\frac{0.9}{1\cdot 0.905}\approx 0.995$)

• 对于第 3 个 token (($t=3$))：

  • ($X_{i,3}=[1,1]$), ($X_{i+1,3}=[0.95,0.95]$)
  • 点积：($X_{i,3}^T{X}_{i+1,3}=(1\cdot 0.95)+(1\cdot 0.95)=0.95+0.95=1.9$)
  • 模长：($\Vert X_{i,3}{\Vert }_2=\sqrt{1^2+1^2}=\sqrt{2}\approx 1.414$), ($\Vert X_{i+1,3}{\Vert }_2=\sqrt{0.95^2+0.95^2}=\sqrt{2\cdot 0.95^2}=\sqrt{2\cdot 0.9025}\approx \sqrt{1.805}\approx 1.343$)
  • 余弦相似度：($\frac{1.9}{1.414\cdot 1.343}\approx \frac{1.9}{1.899}\approx 1.0$)

步骤 2：计算平均余弦相似度 (${\mathbb{E}}_{X,t}$)

假设我们只有这一个样本 ($X$)（在实际中会遍历多个样本），对所有时间步 ($t$) 取平均：
$${\mathbb{E}}_{X,t}\left(\frac{X_{i,t}^T{X}_{i+1,t}}{\Vert X_{i,t}{\Vert }_2\Vert X_{i+1,t}{\Vert }_2}\right)=\frac{0.995+0.995+1.0}{3}\approx 0.9967$$

步骤 3：计算 BI 值

$${\mathrm{BI}}_i=1-{\mathbb{E}}_{X,t}\left(\frac{X_{i,t}^T{X}_{i+1,t}}{\Vert X_{i,t}{\Vert }_2\Vert X_{i+1,t}{\Vert }_2}\right)=1-0.9967=0.0033$$

解释 BI 值：

• BI 值为 0.0033，非常低，说明第 ($i$) 层的输入和输出余弦相似度非常高（接近 1），也就是说这一层对隐藏状态的变换很小，几乎没有改变输入。
• 根据论文的假设，BI 值越低，说明该层越可能是冗余的，可以考虑移除。

---

BI 公式的意义

• 直观理解：BI 衡量的是每一层对隐藏状态的“改造程度”。如果 ($X_{i,t}$) 和 ($X_{i+1,t}$) 几乎相同（余弦相似度接近 1），那么这一层的作用很小，可能是冗余的。
• 实际应用：论文通过 BI 排序所有层，移除 BI 得分最低的层（即最冗余的层），从而实现模型压缩。

---

总结

• ($X_{i,t}$) 是第 ($i$) 层、第 ($t$) 个 token 的隐藏状态向量，但实际计算时 ($X_i$) 是一个矩阵，包含整个序列的隐藏状态。
• BI 公式的核心是计算每一层输入和输出的平均余弦相似度，BI 值越低，说明该层越冗余。
• 通过例子，我们看到 BI 的计算过程：对每个 token 计算余弦相似度，取平均后用 1 减去，得到层的“变换程度”。

后记

2025年5月12日于上海，在grok 3大模型辅助下完成。