Meta
Tag : #articles/tags/Computer_Science__Machine_Learning, #articles/tags//unread, #articles/tags/⭐⭐⭐⭐⭐, #articles/tags/llm/prune, #articles/tags/llm/MoEcitekey : frantarSparseGPTMassiveLanguage2023authors : #articles/authors/Elias-Frantar, #articles/authors/Dan-Alistarhyear : 2023/03/22doi : 10.48550/arXiv.2301.00774url : http://arxiv.org/abs/2301.00774 conference : #articles/conference/Journal : #articles/journals/
Abstract
We show for the first time that large-scale generative pretrained transformer (GPT) family models can be pruned to at least 50% sparsity in one-shot, without any retraining, at minimal loss of accuracy. This is achieved via a new pruning method called SparseGPT, specifically designed to work efficiently and accurately on massive GPT-family models. We can execute SparseGPT on the largest available open-source models, OPT-175B and BLOOM-176B, in under 4.5 hours, and can reach 60% unstructured sparsity with negligible increase in perplexity: remarkably, more than 100 billion weights from these models can be ignored at inference time. SparseGPT generalizes to semi-structured (2:4 and 4:8) patterns, and is compatible with weight quantization approaches. The code is available at: https://github.com/IST-DASLab/sparsegpt .
论文式 (1) 是
arg min mask M ℓ , W ^ ℓ ∥ W ℓ X ℓ − ( M ℓ ⊙ W ^ ℓ ) X ℓ ∥ 2 2 . \arg\min_{\text{mask } M^\ell,\ \hat W^\ell}
\|W^\ell X^\ell - (M^\ell \odot \hat W^\ell)X^\ell\|_2^2.
arg mask M ℓ , W ^ ℓ min ∥ W ℓ X ℓ − ( M ℓ ⊙ W ^ ℓ ) X ℓ ∥ 2 2 .
这里：
W ℓ W^\ell W ℓ ℓ \ell ℓ X ℓ X^\ell X ℓ M ℓ M^\ell M ℓ W ^ ℓ \hat W^\ell W ^ ℓ ⊙ \odot ⊙
M M M 一旦 mask M M M
min W ^ ∥ W X − ( M ⊙ W ^ ) X ∥ F 2 . \min_{\hat W} \|WX - (M\odot \hat W)X\|_F^2.
W ^ min ∥ W X − ( M ⊙ W ^ ) X ∥ F 2 .
这里更准确地写成 Frobenius 范数比较合适，因为输出是矩阵。论文原文写成 ∥ ⋅ ∥ 2 2 \|\cdot\|_2^2 ∥ ⋅ ∥ 2 2
设
W ∈ R d row × d col W \in \mathbb{R}^{d_{\text{row}}\times d_{\text{col}}} W ∈ R d row × d col X ∈ R d col × n X \in \mathbb{R}^{d_{\text{col}}\times n} X ∈ R d col × n
则
W X ∈ R d row × n . WX \in \mathbb{R}^{d_{\text{row}}\times n}.
W X ∈ R d row × n .
把 W W W w i w_i w i i i i
w i X ∈ R 1 × n . w_i X \in \mathbb{R}^{1\times n}.
w i X ∈ R 1 × n .
同理，压缩后第 i i i
( M i ⊙ w ^ i ) X . (M_i\odot \hat w_i)X.
( M i ⊙ w ^ i ) X .
所以整体目标可以拆成
∥ W X − ( M ⊙ W ^ ) X ∥ F 2 = ∑ i = 1 d row ∥ w i X − ( M i ⊙ w ^ i ) X ∥ 2 2 . \|WX-(M\odot \hat W)X\|_F^2
=
\sum_{i=1}^{d_{\text{row}}}
\|w_iX-(M_i\odot \hat w_i)X\|_2^2.
∥ W X − ( M ⊙ W ^ ) X ∥ F 2 = i = 1 ∑ d row ∥ w i X − ( M i ⊙ w ^ i ) X ∥ 2 2 .
结论：固定 mask 后，每一行可以独立求解。
对应 mask：M i ∈ { 0 , 1 } 1 × d col M_i \in \{0,1\}^{1\times d_{\text{col}}} M i ∈ { 0 , 1 } 1 × d col
设保留下来的列索引集合为
S i = { j : M i j = 1 } . S_i=\{j: M_{ij}=1\}.
S i = { j : M i j = 1 } .
把这些保留位置抽出来：
w i , S i w_{i,S_i} w i , S i w i w_i w i X S i X_{S_i} X S i X X X
那么因为被剪掉的位置必须为 0，所以优化变量只剩保留位置上的权重。记它为 u u u
min u ∥ w i X − u X S i ∥ 2 2 . \min_u \|w_iX-uX_{S_i}\|_2^2.
u min ∥ w i X − u X S i ∥ 2 2 .
这就是一个标准线性最小二乘。
现在我们来严格推导闭式解。
令
y i : = w i X ∈ R 1 × n , y_i := w_iX \in \mathbb{R}^{1\times n},
y i : = w i X ∈ R 1 × n ,
则问题是
min u ∥ y i − u X S i ∥ 2 2 . \min_u \|y_i-uX_{S_i}\|_2^2.
u min ∥ y i − u X S i ∥ 2 2 .
把它写成转置更接近标准回归形式也可以：
min u ⊤ ∥ y i ⊤ − X S i ⊤ u ⊤ ∥ 2 2 . \min_{u^\top} \|y_i^\top - X_{S_i}^\top u^\top\|_2^2.
u ⊤ min ∥ y i ⊤ − X S i ⊤ u ⊤ ∥ 2 2 .
f ( u ) = ∥ y i − u X S i ∥ 2 2 = ( y i − u X S i ) ( y i − u X S i ) ⊤ . f(u)=\|y_i-uX_{S_i}\|_2^2
=(y_i-uX_{S_i})(y_i-uX_{S_i})^\top.
f ( u ) = ∥ y i − u X S i ∥ 2 2 = ( y i − u X S i ) ( y i − u X S i ) ⊤ .
展开：
f ( u ) = y i y i ⊤ − 2 y i X S i ⊤ u ⊤ + u X S i X S i ⊤ u ⊤ . f(u)=y_iy_i^\top -2y_iX_{S_i}^\top u^\top + uX_{S_i}X_{S_i}^\top u^\top.
f ( u ) = y i y i ⊤ − 2 y i X S i ⊤ u ⊤ + u X S i X S i ⊤ u ⊤ .
u u u 对行向量形式求导，得到正规方程：
u X S i X S i ⊤ = y i X S i ⊤ . uX_{S_i}X_{S_i}^\top = y_iX_{S_i}^\top.
u X S i X S i ⊤ = y i X S i ⊤ .
也就是
u H S i = y i X S i ⊤ , 其中 H S i : = X S i X S i ⊤ . uH_{S_i}=y_iX_{S_i}^\top,
\quad \text{其中 } H_{S_i}:=X_{S_i}X_{S_i}^\top.
u H S i = y i X S i ⊤ , 其中 H S i : = X S i X S i ⊤ .
若 H S i H_{S_i} H S i
u = y i X S i ⊤ ( X S i X S i ⊤ ) − 1 . u = y_iX_{S_i}^\top (X_{S_i}X_{S_i}^\top)^{-1}.
u = y i X S i ⊤ ( X S i X S i ⊤ ) − 1 .
再代入 y i = w i X y_i=w_iX y i = w i X
u = ( w i X ) X S i ⊤ ( X S i X S i ⊤ ) − 1 . u = (w_iX)X_{S_i}^\top (X_{S_i}X_{S_i}^\top)^{-1}.
u = ( w i X ) X S i ⊤ ( X S i X S i ⊤ ) − 1 .
这就是标准最小二乘解。
这是一个实际问题，论文在算法部分也用了 dampening：
H − 1 = ( X X ⊤ + λ I ) − 1 H^{-1}=(XX^\top + \lambda I)^{-1}
H − 1 = ( X X ⊤ + λ I ) − 1
说明他们也默认实际中会做正则化稳定求逆。
所以严格来说，如果
X S i X S i ⊤ X_{S_i}X_{S_i}^\top
X S i X S i ⊤
不可逆或病态，常见做法是解 ridge 版本：
u ∗ = y i X S i ⊤ ( X S i X S i ⊤ + λ I ) − 1 . u^* = y_iX_{S_i}^\top (X_{S_i}X_{S_i}^\top+\lambda I)^{-1}.
u ∗ = y i X S i ⊤ ( X S i X S i ⊤ + λ I ) − 1 .
这个思想和 SparseGPT 后面使用 Hessian dampening 是一致的。
聪明的做法是：
把 w m w_m w m
同时调一下别的权重，让整体输出尽量维持不变
在“第 m m m 并且由于直接求解闭式解非常昂贵, 与其直接对完整 mask M M M
下面进入核心。
设当前一行权重 是列向量 w ∈ R d w\in\mathbb{R}^d w ∈ R d
假设当前 w w w w w w L ( w ) L(w) L ( w )
L ( w + δ ) ≈ L ( w ) + 1 2 δ ⊤ H δ L(w+\delta) \approx L(w) + \frac12 \delta^\top H \delta
L ( w + δ ) ≈ L ( w ) + 2 1 δ ⊤ H δ
其中：
H H H 因为当前点梯度为 0，所以一阶项没了
如果要删掉第 m m m
w m + δ m = 0 w_m + \delta_m = 0
w m + δ m = 0
注意这里左边的 δ m \delta_m δ m m m m Δ \Delta Δ
所以约束是：
e m ⊤ Δ = − w m e_m^\top \Delta = -w_m
e m ⊤ Δ = − w m
其中 e m e_m e m m m m
于是问题变成：
min Δ 1 2 Δ ⊤ H Δ s.t. e m ⊤ Δ = − w m \min_\Delta \frac12 \Delta^\top H \Delta
\quad
\text{s.t. } e_m^\top \Delta = -w_m
Δ min 2 1 Δ ⊤ H Δ s.t. e m ⊤ Δ = − w m
这就是 OBS 的本质优化问题：
最小化二次损失增加，同时强制第 m m m
这是最关键的部分。
L ( Δ , λ ) = 1 2 Δ ⊤ H Δ + λ ( e m ⊤ Δ + w m ) \mathcal{L}(\Delta,\lambda)
=
\frac12 \Delta^\top H \Delta + \lambda (e_m^\top \Delta + w_m)
L ( Δ , λ ) = 2 1 Δ ⊤ H Δ + λ ( e m ⊤ Δ + w m )
Δ \Delta Δ ∇ Δ L = H Δ + λ e m = 0 \nabla_\Delta \mathcal{L} = H\Delta + \lambda e_m = 0
∇ Δ L = H Δ + λ e m = 0
所以
Δ = − λ H − 1 e m \Delta = -\lambda H^{-1} e_m
Δ = − λ H − 1 e m
而 H − 1 e m H^{-1}e_m H − 1 e m H − 1 H^{-1} H − 1 m m m
Δ = − λ H : , m − 1 \Delta = -\lambda H^{-1}_{:,m}
Δ = − λ H : , m − 1
λ \lambda λ 约束要求：
e m ⊤ Δ = − w m e_m^\top \Delta = -w_m
e m ⊤ Δ = − w m
代入上式：
e m ⊤ ( − λ H : , m − 1 ) = − w m e_m^\top(-\lambda H^{-1}_{:,m}) = -w_m
e m ⊤ ( − λ H : , m − 1 ) = − w m
即
− λ [ H − 1 ] m m = − w m -\lambda [H^{-1}]_{mm} = -w_m
− λ [ H − 1 ] m m = − w m
所以
λ = w m [ H − 1 ] m m \lambda = \frac{w_m}{[H^{-1}]_{mm}}
λ = [ H − 1 ] m m w m
Δ = − w m [ H − 1 ] m m H : , m − 1 \Delta
=
-\frac{w_m}{[H^{-1}]_{mm}} H^{-1}_{:,m}
Δ = − [ H − 1 ] m m w m H : , m − 1
这就是论文式 (3)：
δ m = − w m [ H − 1 ] m m H : , m − 1 \delta_m = -\frac{w_m}{[H^{-1}]_{mm}} H^{-1}_{:,m}
δ m = − [ H − 1 ] m m w m H : , m − 1
除此之外我们还可以得到:
ε m = w m 2 [ H − 1 ] m m \varepsilon_m = \frac{w_m^2}{[H^{-1}]_{mm}}
ε m = [ H − 1 ] m m w m 2
删除某个权重是否“容易”，不只取决于它绝对值小不小，还取决于：
[ H − 1 ] m m [H^{-1}]_{mm}
[ H − 1 ] m m
如果这个值大，说明这个权重的影响更容易被其他参数补偿，那么删除代价就可能更小。这在后续会用来自适应mask
H − 1 H^{-1} H − 1 H H H H − 1 H^{-1} H − 1
这个部分是整篇文章最难的理解的部分
事实上并不能证明这样做和原来的OBS求解保持等价
每隔B s B_s B s ε m \varepsilon_m ε m
step freeze wights 是因为这里是mask的是需要保留的.
所以这里E E E w c 2 [ H − 1 ] c c 2 \frac{w_c^2}{[H^{-1}]^2_{cc}} [ H − 1 ] c c 2 w c 2
外层每隔B B B
1 2 3 4 5 6 7 8 9 10 11 12 13 矩阵 W 被分为多个 Block (大小为 B) \[ 已剪枝固定区 \] | \[ 当前正在处理的 Block B \] | \[ 尚未处理的剩余庞大区域 (i+B 之后) \] | (列索引 i 到 i+B-1) | | | | ^ 当前列 j | | |--内层局部 Update---> | | (仅更新 j 到 i+B 的列) | | | 当整个 Block B 的循环走完，累积了误差矩阵 E: |-----------------------外层全局 Update--------------------------->| | (E * H^{-1} 一次性矩阵乘法更新右侧所有列) |
本篇文章的思路实际非常清晰明了, 作者的整个逻辑链条 大概是:
将非结构化剪枝问题建模成一个优化问题
尝试求解这个优化问题的闭式解 ==> 求逆的代价很高
使用OBS中迭代地进行处理这个问题 ==> 为什么不使用最优部分更新? 来减少每次迭代导致的大范围参数更新
事实上每一步都是由上一步产生的问题而引申出来的解决办法.
takeaway
X 1 ∈ R m × n X_1\in\mathbb{R}^{m\times n} X 1 ∈ R m × n X 2 ∈ R n × p X_2\in\mathbb{R}^{n\times p} X 2 ∈ R n × p O ( m n p ) \mathcal{O}(mnp) O ( m n p ) O ( n 3 ) \mathcal{O}(n^3) O ( n 3 )
对于样本 X ∈ R B × N × d i n X\in\mathbb{R}^{B\times N\times d_{in}} X ∈ R B × N × d i n
1 nn.Linear(input_dim, output_dim)
那么：
weight.shape == (output_dim, input_dim)bias.shape == (output_dim,)
前向实际上做的是：
y = X W ⊤ + b y = XW^\top + b
y = X W ⊤ + b
所以虽然概念上你也可以说“把输入乘一个 input_dim 到 output_dim 的映射矩阵”，但参数真正存储的 shape 是：
( output_dim , input_dim ) (\text{output\_dim},\text{input\_dim})
( output_dim , input_dim )
并且非常重要的一点是, 从矩阵乘法的角度来看我们可以知道权重矩阵的哪些参数和哪些样本有关 , 对于去掉Batch的样本矩阵, 可以按照列分块的形式写成向量( x 1 , x 2 , x 3 . . . ) (x_1,x_2,x_3...) ( x 1 , x 2 , x 3 . . . ) ( w 1 , w 2 . . . . ) (w_1,w_2....) ( w 1 , w 2 . . . . ) input neuron id .
