损伤感知多臂老虎机剪枝法提升模型压缩效果
「研究人员提出一种名为“损伤感知多臂老虎机剪枝”的新方法,用于视觉与语言Transformer模型的结构化后训练剪枝。该方法通过评估临时掩蔽功能单元造成的损失变化,在固定评估预算下优化剪枝选择,显著降低模型性能退化。」
近日,一篇发表于arXiv的论文《Damage-Aware Bandit Pruning for Vision and Language Transformers》提出了一种创新的结构化剪枝策略,旨在在不显著损害模型性能的前提下,高效地压缩大型Transformer模型。该方法将剪枝问题建模为一个“损伤感知的多臂老虎机”(damage-aware multi-armed bandit)问题,在有限的候选单元评估预算内,智能地选择对模型整体性能影响最小的功能单元进行移除。
传统剪枝方法如随机剪枝、基于权重幅度的剪枝或静态显著性剪枝,往往无法充分考虑不同结构单元之间的相互依赖关系,容易导致性能大幅下降。而本文提出的方法则聚焦于结构化单元——例如注意力头(attention heads)和MLP通道组(MLP channel groups)——在小批量校准数据上进行临时掩蔽,并计算“配对损伤”(paired damage),即掩蔽后的损失与原始损失之差。这一设计有效减少了批次间的波动,使评估更加稳定可靠。
在此基础上,研究团队设计了一种平滑且有界的奖励函数,用于驱动两种不同的探索策略:一种是基于UCB(Upper Confidence Bound)风格的策略,另一种则是分数Beta分布的汤普森采样(fractional-Beta Thompson Sampling)。最终的剪枝掩码通过逐步添加被选中的单元构建而成,确保每一步都基于当前最优信息进行决策。
值得注意的是,该方法并不直接追求物理层面的模型压缩或推理加速,而是关注于“有效结构抑制”——即将选定的单元在原始稠密模型检查点中功能归零,从而真实反映剪枝对模型能力的影响。实验覆盖了多种主流语言模型(如GPT-2、OPT、Pythia、Qwen2.5、SmolLM2)和视觉模型(如ViT-B/16、DeiT-Tiny、Swin-Tiny),在WikiText-2、LAMBADA和Imagenette等基准数据集上进行了全面评估。
结果表明,在五次随机种子实验中,该带宽控制下的老虎机方法在多数语言模型对比中优于传统的预算贪婪剪枝策略。论文重点报告的28组对比中,有23组的自助置信区间排除了零效应,其中11组的配对检验p值小于0.05;经过Benjamini-Hochberg多重检验校正后,仍有6组在全部116项测试中保持q值低于0.05,显示出统计显著性。此外,针对ViT-B/16和Swin-Tiny的匹配评估进一步证实,性能提升并非单纯源于更大的评估预算,而是方法本身的优越性所致。
这项工作为后训练结构化剪枝提供了一个兼顾效率与性能的新范式,尤其适用于资源受限但需保持高精度的AI部署场景,有望推动轻量化大模型在边缘设备和实际应用中的落地。
来源:Heooo AI工具导航