MOONSHOT:Google与OpenAI提出多目标剪枝框架,大模型困惑度降低32.6%

MOONSHOT : A Framework for Multi-Objective Pruning of Vision and Large Language Models

MOONSHOT:Google与OpenAI提出多目标剪枝框架,大模型困惑度降低32.6% 论文图示

随着当代视觉与大语言模型(LLMs)参数规模的指数级增长,其在推理阶段所消耗的计算资源与显存带宽已成为实际部署中的严峻挑战。在众多模型压缩策略中,网络剪枝(Network Pruning)因其能够直接移除冗余权重而备受关注。特别是在大模型时代,传统的“剪枝-微调”范式由于高昂的重训练成本变得不再现实。因此,研究界的焦点逐渐转向了极具挑战性的后训练单次剪枝(post-training one-shot pruning)——仅需少量校准数据,且无需任何重新训练即可完成模型压缩。

ArXiv URL:https://arxiv.org/abs/2604.13287v1

然而,现有的单次剪枝算法通常受限于单一的优化目标。它们要么侧重于最小化局部特征的逐层重构误差,要么侧重于全局训练损失的二阶泰勒近似。来自 Google、MIT 以及 OpenAI 的研究团队发现,这两种截然不同的剪枝标准在不同的网络架构和稀疏度层级下互有胜负,单一目标往往会导致次优的剪枝决策。

基于这一关键洞察,研究人员提出了 MOONSHOT(Multi-Objective ONe-SHOT pruning)框架。这是一种通用且高度灵活的多目标优化框架,通过联合优化逐层重构误差与训练损失的二阶泰勒近似,巧妙地整合了局部与全局的参数重要性信号。作为一个强大的“包装器”(wrapper),MOONSHOT 可以无缝集成到现有主流剪枝算法中。为了在十亿乃至百亿参数规模上保持极高的运算效率,作者团队设计了一套近似逆 Hessian 矩阵的高效计算流程。

实验结果证明了该框架的卓越有效性:在 Llama-3.2 和 Llama-2 模型上,将 MOONSHOT 与最先进的剪枝方法结合,在 2:4 半结构化稀疏下可使 C4 数据集的困惑度大幅降低多达 32.6%,并在七个分类基准测试中将零样本平均准确率提升 4.9 个点。而在视觉模型领域,MOONSHOT 同样使 Vision Transformers 在 70% 稀疏度下的 ImageNet-1k 准确率提升超过 5 个点,ResNet-50 在 90% 稀疏度下提升 4 个点。

单一剪枝目标的局限性与互补性信号

要理解 MOONSHOT 的创新逻辑,首先需要深入探究现有后训练单次剪枝方法所依赖的基础评价标准。在确定哪些权重应当被保留、哪些应当被归零时,算法需要一种量化指标来估计权重移除对模型最终性能的损害。除了最基础的基于权重绝对值大小(Magnitude)的策略外,当前最先进的方法主要分为两大阵营:

第一类阵营侧重于训练损失的局部二次近似。这类方法(如 CAP 算法)追溯至经典的 Optimal Brain Surgeon (OBS) 框架。其核心思想是利用原始网络训练损失函数的二阶泰勒展开,通过计算目标函数关于权重的 Hessian 矩阵,来全局评估每个参数对整体训练误差的影响。这种方法利用了网络在训练阶段学到的全局梯度特征。

第二类阵营则采用逐层重构误差(Layer-wise Reconstruction Error)。以 SparseGPT 和 Wanda 等专为大语言模型设计的明星算法为代表,这类方法将庞大的网络剪枝问题拆解为多个独立的单层子问题。对于每一层,算法的目标是在给定的稀疏度约束下,使剪枝后该层输出的激活值尽可能接近原始未剪枝前的输出。重构误差更聚焦于局部嵌入空间的信息保真度。

研究团队对这两种经典标准进行了系统性对比。他们让主要针对视觉架构设计的 CAP 算法去优化逐层重构误差,同时让基于重构误差设计的 OBC、SparseGPT 和 Wanda 算法去优化训练损失的二阶近似。结果揭示了一个重要现象:没有任何单一的优化目标能够在所有模型架构和稀疏度水平下保持绝对的领先。

例如,在某些特定的视觉特征提取层或极高稀疏度下,基于全局训练损失近似的准则能够更好地维持网络功能;而在大语言模型的多头注意力层中,局部激活值的保真度往往更能直接反映生成质量。这种性能的交替领先强烈暗示了一个结论:这两种目标分别捕获了关于网络参数重要性的互补信号。仅依赖其中之一,不可避免地会遗漏关键维度的信息,从而产生次优的掩码分配。

MOONSHOT 的核心机制:双重目标的融合与数学重构

既然局部重构与全局近似包含互补信息,最直观的改进方向便是将两者结合。MOONSHOT 框架正是建立在这一逻辑之上,它通过一个权重参数将这两个目标融合为一个全新的多目标代价函数。

为了确保两个量级可能完全不同的目标函数能够被平等且稳定地优化,MOONSHOT 在组合它们之前引入了规范化操作。具体而言,算法分别用两个目标函数在全零权重矩阵(即极端情况下的完全剪枝)下的取值作为分母进行缩放。设定一个调节因子 $\lambda \in [0,1]$,合并后的目标函数被定义为:

\[\mathcal{L}^{(l)}_{\lambda} := ({\lambda}/{\mathcal{L}^{(l)}_{R}(\mathbf{0})})\mathcal{L}^{(l)}_{R} + ((1-\lambda)/{\mathcal{L}^{(l)}_{F}(\mathbf{0})})\mathcal{L}^{(l)}_{F}\]

其中,$\mathcal{L}^{(l)}{R}$ 代表第 $l$ 层的局部重构误差,而 $\mathcal{L}^{(l)}{F}$ 则代表基于 Fisher 信息的二阶泰勒近似损失。

通过这种重新公式化,多目标优化问题依然可以被转化为一个带有基数约束(即稀疏度约束)的凸二次规划问题。从理论层面上看,MOONSHOT 将目标函数的二次型重新整理,使得问题形式再次与单一目标下的标准格式对齐。这意味着,理论上现存的所有基于 OBS 规则的单目标求解器,都可以直接接管并优化这个新的多目标函数。

然而,在将这套优美的理论向十亿级以上参数的大型模型迁移时,工程层面的算力噩梦随之浮现。

MOONSHOT:Google与OpenAI提出多目标剪枝框架,大模型困惑度降低32.6% 论文图示

上图展示了不同目标下 Hessian 矩阵的结构差异。在单目标逐层重构误差的设定下(如最左侧所示),输入特征带来的 Hessian 矩阵具有极其完美的块对角(Block-diagonal)结构,且更重要的是,对角线上的每一个数据块都是完全相同的。这使得底层算法(如 SparseGPT)在处理包含成千上万行的权重矩阵时,只需计算一次相对较小矩阵的逆,然后复用于所有行,极大地降低了计算复杂度。

但在引入了包含二阶泰勒近似的多目标融合后(如最右侧所示),即便研究人员强行保留块对角近似,对角线上的各个数据块也不再相同。此时,针对不同的权重行,算法需要独立面对不同且庞大的 Hessian 块。如果采用传统的 Cholesky 分解进行直接求逆,所需的运算时间与显存将随着模型维度的增加而呈立方级爆炸,这直接摧毁了单次剪枝“轻量、快速”的核心优势。

跨越算力瓶颈:大规模模型下的高效逆 Hessian 计算

为了让 MOONSHOT 能够真正应用在现代大语言模型(如 Llama-3.2)上,作者团队提出了一套针对融合 Hessian 矩阵求逆的高效近似程序。

核心的数学武器是 Woodbury 矩阵恒等式(Woodbury Matrix Identity)。由于新的 Hessian 矩阵构成的特点可以拆解为一个基准矩阵加上一个由特定维度控制的低秩微调项,直接对目标矩阵求逆的繁重任务,被等效转换为了对一个小规模矩阵求逆的轻量级操作。

在 MOONSHOT 的具体实现中,对于语言模型的某一层,计算规模被巧妙地压缩到了一定维度的特征空间内。通过这种重构,系统不再需要对每个不再相同的大尺寸分块进行全矩阵求逆,而是只需进行少量的基础矩阵乘法和极低维度的求逆运算。据论文材料披露,在处理 Llama-3.2-3B 模型规模最大的层时,这一高效逆运算的执行时间仅需最多 5 到 6 秒钟,相较于直接调用标准 Cholesky 分解矩阵求逆,实现了高达 6 倍的提速。

正是因为这套底层计算逻辑的优化,MOONSHOT 作为底层基建,顺利承载起了庞大的算力压力。最终,MOONSHOT 包装后的 SparseGPT 算法在单张普通专业级 GPU 上,可以在 40 分钟内完成对 Llama-3.2-3B 模型的全局高质量剪枝;对于 1B 模型更是仅需 8 分钟即可完成。多目标优化的引入并未以牺牲单次剪枝的实用效率为代价。

实验验证:多维度稀疏场景下的全面超越

为了验证框架的稳健性与通用性,研究团队在不同类型的数据集、不同的网络架构以及各种稀疏化场景下,对 MOONSHOT 进行了严苛的评估。现代深度学习应用中的稀疏模式主要分为三类,MOONSHOT 对其提供了全覆盖的优化:

  1. 无结构稀疏(Unstructured Sparsity):随意剪除特定比例的权重。虽然能极大地缩小模型体积并在专有硬件上提速,但由于缺乏规律,往往需要极高的稀疏度才能在通用 GPU 上见到成效,这通常会严重破坏模型性能。

  2. 结构化稀疏(Structured Sparsity):直接剪除整个通道或权重行。能提供立竿见影的加速,但模型能力衰减极快,仅能在低稀疏度下使用。

  3. 半结构化稀疏(Semi-structured n:m Sparsity):在每 m 个连续权重中强制归零 n 个(通常为 2:4)。这是目前能够完美匹配 NVIDIA Ampere 等现代 GPU 架构硬件加速、且能较好保持模型质量的最佳折中方案。

MOONSHOT:Google与OpenAI提出多目标剪枝框架,大模型困惑度降低32.6% 论文图示

从上文的综合性能走势图中可以清晰地看到 MOONSHOT 带来的绝对增益。

在大型语言模型(LLMs)上的表现:

研究团队将 MOONSHOT 与大语言模型界最知名的剪枝基线 SparseGPT 和 Wanda 进行了结合。在 Llama-3.2(1B 和 3B 版本)以及广泛使用的 Llama-2-13b-chat 模型上,多目标优化的威力显露无疑。特别是在工业界最为青睐的 2:4 稀疏度场景下,MOONSHOT 成功将 C4 文本测试集上的困惑度(Perplexity)降低了最高达 32.6%。困惑度是衡量语言模型预测下一个词能力的最核心指标,超过 30% 的相对下降意味着模型在相同剪枝约束下,语言连贯性和知识保留度实现了质的飞跃。

同时,在常识推理与分类能力方面,经过 MOONSHOT 增强的模型在包括 PIQA、HellaSwag、ARC 等七大零样本(Zero-shot)评测基准上的平均准确率提升了最多 4.9 个点。即使在 10% 结构化稀疏这种极易导致模型崩溃的设定下,MOONSHOT 协助 OSSCAR 算法依然降低了 11% 的困惑度。

更为关键的是,MOONSHOT 的优化具有正交性。近年来,许多研究致力于通过非均匀稀疏分配(如 OWL 和 AlphaPruning 算法)在各层间动态调配剪枝比例来保住大模型性能。实验证明,MOONSHOT 能够与这些非均匀分配策略完美兼容。在非均匀分配已带来巨大提升的基础之上,MOONSHOT 进一步挖掘了局部权重的保留价值,在极限压缩状态下依然能榨取额外约 25% 的困惑度下降。

在视觉前馈网络(Vision Models)上的表现:

虽然目前关注焦点常集中在大模型,但 MOONSHOT 在视觉任务上的表现同样强劲。在基于 Transformer 架构的 DeiT 序列中,结合了 MOONSHOT 的 CAP 算法在 70% 的高稀疏度下,将 ImageNet-1k 上的分类准确率直接拔高了超过 5 个百分点。而在经典的卷积神经网络 ResNet-50 上,在极端严苛的 90% 稀疏度下,MOONSHOT 使 OBC 算法实现了 4 个百分点的性能回升。这些数据确凿地证明,融合重构损失与局部泰勒近似的多目标逻辑,不受限于特定的自注意力或卷积操作机制,具有广泛的网络拓扑适用性。

总结与展望

在模型越发庞大、算力愈发昂贵的当下,如何“优雅地缩小”模型已经成为了决定 AI 技术商业落地可行性的核心议题。传统的单次后训练剪枝算法为了追求计算速度,在目标函数的设计上往往做出妥协,这在参数规模相对较小的时代尚可接受,但在动辄数十亿参数的 LLM 中,粗放的局部优化会带来不可逆的知识遗忘。

MOONSHOT 框架的提出,标志着模型剪枝技术对优化目标本质理解的一次重要升级。它用严密的对照实验揭露了“重构误差”与“二阶泰勒近似”各自的盲区,并运用数学重组与 Woodbury 快速求解技术,巧妙地将这两股原本平行的力量汇聚到了一起。

从单一目标到多维信息融合的跨越,不但为现有的先进剪枝工具套上了一层强力的增幅装甲,更揭示了压缩领域未来的可能路径:如果局部特征重构与全局损失近似能够如此成功地互补,未来是否可以进一步引入更多维度的激活分布特征或是注意力机制的特定导向作为优化信号?MOONSHOT 已经为这条更精细、更多元的大模型压缩之路奠定了坚实的基石。