Meta

Tag: #articles/tags//unread
citekey: leeSTUNStructuredThenUnstructuredPruning
authors: #articles/authors/Jaeseong-Lee, #articles/authors/Seung-won-Hwang, #articles/authors/Aurick-Qiao, #articles/authors/Daniel-Campos, #articles/authors/Zhewei-Yao, #articles/authors/Yuxiong-He
year: Error: format can only be applied to dates. Tried for format object
doi:
url:
conference: #articles/conference/
Journal : #articles/journals/

Abstract

Mixture-of-experts (MoEs) have been adopted to reduce inference costs by sparsely activating experts in large language models (LLMs). Despite these reductions, the massive number of parameters in MoEs still makes them expensive to serve. Conventionally, unstructured or structured pruning has been considered to reduce the number of parameters. Our key contribution is exploring the interpolation between structured and unstructured pruning, to propose a novel structured-then-unstructured (STUN) approach outperforming both structured and unstructured pruning, especially for MoEs. In the first stage, we show a scalable expert pruning with O(1) forward pass, unlike existing work requiring O(knn)O(\frac{k^n}{\sqrt{n}}) forward passes for n experts that cannot scale for recent MoEs with hundreds of experts. We then show our expert-pruned MoEs are robust to unstructured pruning to follow. Experiments on Snowflake Arctic and Mixtral show that our proposal is highly effective– For Snowflake Arctic, a 480B-sized MoE with 128 experts, our method needs only one H100 and two hours to achieve nearly no loss in performance with 40% sparsity, even in generative tasks such as GSM8K, where state-of-the-art structured or unstructured pruning methods fail. The code is publicly available.

Note

1. how we remove redundant experts with expert-level structured pruning with high scalability

1.1 优化目标

ES=∥M(x;θ)−M(x;θ−θS)∥F\mathcal{E}_S = \|M(x; \theta) - M(x; \theta - \theta_S)\|_F

由于这个优化问题非常的难解(需要遍历所有的组合方式), 于是作者想到使用贪心的方法来减少计算量:

1.2 原问题的精确解

精确解需要遍历所有的专家剪枝的组合方式, 我们转换为一个概率问题, 可以形式化为如下的形式:

argmaxSP(X1=s1,⋯ ,X∣S∣=s∣S∣)\text{argmax}_S P(X_1 = s_1, \cdots, X_{|S|} = s_{|S|})

P(X1=s1,⋯ ,X∣S∣=s∣S∣)=1Z⋅1ES(10)P(X_1 = s_1, \cdots, X_{|S|} = s_{|S|}) = \frac{1}{Z} \cdot \frac{1}{\mathcal{E}_S} \quad (10)

1.3. 简化和贪心

  • 对于上面这个过程我们可以转换为逐步求解和优化的过程, 但是这其中做了一些简化, 比如剪枝某个专家的概率是不变的并且等于只剪枝该专家的概率, 剪枝某个专家的概率与其他(不属于同一组的)专家是无关的.

1.3.1. 专家分组策略

  • 专家行为相似度 (Behavioral Similarity) 无需运行大量数据，仅通过观察路由器权重 WW，计算其矩阵距离。距离越小，这两个专家处理的任务越相似 。$$b_{i,j} = -\lambda_{1} ||W_{i} - W_{j}||{F} + \lambda{2} a_{i,j}$$
    • 🔍 Variable / Detail Supplement: bi,jb_{i,j}: 专家 ii 和专家 jj 之间的行为相似度 。 ∣∣Wi−Wj∣∣F||W_{i} - W_{j}||_{F}: 路由器权重矩阵行向量之间的 Frobenius 范数（可以理解为欧氏距离的矩阵版） 。 ai,ja_{i,j}: 共激活统计频率（可选使用，通过极少量校准数据获得） 。 λ1,λ2\lambda_{1}, \lambda_{2}: 平衡两种相似度信号的超参数 。

258

1.3.2. 贪心: 条件概率下的极大似然估计

  • 其实本质上作者将: 在不同情况下, 剪枝某个专家的代价设为只剪枝改专家的代价, 并通过一定的启发式进行搜索, 至于能不能搜索到最优解, 存疑
  • 通过这种方式可以将剪枝的复杂度从枚举所有的情况转化为只需要求解只剪枝每个专家的代价即可, 复杂度是n(专家数)次forward.

arg⁡maxS∏kP(Xk=sk∣X1=s1,⋅⋅⋅,Xk−1=sk−1)\arg max _S ∏^k P (X_k = s_k|X_1 = s_1, · · · , X_{k−1} = s_{k−1})

P(Ei∣Sk−1)={P(Ei)−pc(Ei)⊆SkP(Ei)otherwiseP(E_i | S_{k-1}) = \begin{cases} P(E_i) - p & c(E_i) \subseteq S_k \\ P(E_i) & \text{otherwise} \end{cases}

1.4. 一阶近似简化

作者继续优化, 将优化目标使用一阶近似, 这样可以达到O(1)\mathcal{O}(1)的forward复杂度. 如果使用某个参数θC\theta_C替换当前专家, 通过对优化目标进行一阶泰勒展开得:

Ei=∣∣E′(θi)⋅(θi−θC)∣∣2\mathcal{E}_{i} = ||E^{\prime}(\theta_{i}) \cdot (\theta_{i} - \theta_{C})||^{2}

通常认为模型是训练的, 即达到最优点一阶导数很小, 则该误差的上限受制于 ∣∣θi−θC∣∣2||\theta_{i} - \theta_{C}||^{2}。因此，数学结论非常直观：保留那个最接近整个聚类平均权重 θi‾\overline{\theta_{i}} 的专家 θC\theta_{C}，就能让整体误差最小化, 让我联想到了VQVAE里的codebook

333

防偏离机制：选择性重建 (Selective Reconstruction)：

  • ⚠️ Limitation/Note: 有时候，直接用平均权重代替保留下来的专家权重能进一步降低当前层的重构误差 。但是，这可能会导致输出分布发生改变，坑了后面一层的网络（Distribution Shift） 。 因此，作者加了一个判断：当总剩余簇数量 ∣A∣|\mathcal{A}| 小于阈值 κ\kappa 时才启用平均权重合并，否则老老实实保留原本存在的那个专家 。

2. how we perform unstructured pruning inside individual experts

专家级冗余扫清后，调用现成的非结构化剪枝算法（如 OWL），对剩下的王牌专家进行细粒度的权重掩码修剪（剔除接近0的权重） 。因为第一阶段没伤筋动骨（峰度极高），第二阶段下刀就可以非常稳 。

Summary

Q&A

  1. 这篇论文的稀疏度很迷,
    1. 非常难以精准的调整, 因为这个稀疏度是与聚类相关的, 只能通过一个超参数来调整不能直接设置.
    2. 其次就是, 这里的稀疏度包含了专家drop/merge + 非结构化剪枝的稀疏, 作者并没有阐述在相同的稀疏度下这两者稀疏度之间的关系
  2. 其次, 就是数学推导不严谨, 做了很多过于松弛的假设, 个人认为还不如不提数学, 而直接从intuition的角度出发解释.
  3. 在非结构化剪枝的部分, 作者几乎没有仔细说他们是怎么做的.
+ takeaway

what make a good prune

  1. We support STUN with the findings of MasonWilliams and Dahlqvist (2024), which show that higher kurtosis in the weight distribution (indicating many outliers) suggests more weights can be pruned while maintaining performance, highlighting the robustness of unstructured pruning. We argue that expert-level pruning does not reduce kurtosis, thereby preserving the network’s resilience to unstructured pruning.
  2. 到时候我们可以看看不同模型的这个kurtosis和我们的最终剪枝结果有什么联系. 或者进一步探索这其中的insight也许会有些帮助