Damage-Aware Bandit Pruning for Vision and Language Transformers
Damage-Aware Bandit Pruning for Vision and Language Transformers
面向视觉与语言 Transformer 的“损伤感知”多臂老虎机剪枝方法
Abstract: Structured post-training pruning of transformers requires selecting complete functional units whose suppression causes limited degradation. We formulate structured-unit selection for language and vision transformers as a damage-aware multi-armed bandit problem under a fixed candidate-evaluation budget.
摘要： Transformer 的结构化训练后剪枝需要筛选出那些在被抑制后只会导致有限性能下降的完整功能单元。我们将语言和视觉 Transformer 的结构化单元选择问题，建模为在固定候选评估预算下的“损伤感知”多臂老虎机（Multi-armed Bandit）问题。
Attention heads and MLP channel groups are temporarily masked on calibration batches. Paired damage is the masked loss minus the base loss on the same batch, reducing batch-to-batch variation. A smooth bounded reward drives either a UCB-style policy or fractional-Beta Thompson Sampling, and the final mask is constructed sequentially by adding one unit at each step.
注意力头和 MLP 通道组会在校准批次上被临时掩码。配对损伤（Paired damage）定义为掩码后的损失减去同一批次上的基础损失，从而减少了批次间的差异。平滑的有界奖励机制驱动着 UCB 风格的策略或分数 Beta 汤普森采样（Fractional-Beta Thompson Sampling），最终的掩码通过每步增加一个单元的方式顺序构建。
The selected units are functionally zeroed in the original dense checkpoint; therefore, the reported parameter effects represent effective structural suppression rather than physical compression or measured speedup.
所选单元在原始稠密检查点中被功能性置零；因此，报告的参数影响代表的是有效的结构抑制，而非物理压缩或实际测得的加速。
Experiments on WikiText-2, LAMBADA, and Imagenette cover GPT-2, OPT, Pythia, Qwen2.5, SmolLM2, ViT-B/16, DeiT-Tiny, and Swin-Tiny, with comparisons against random, magnitude, static-saliency, and budgeted-greedy selection.
在 WikiText-2、LAMBADA 和 Imagenette 数据集上的实验涵盖了 GPT-2、OPT、Pythia、Qwen2.5、SmolLM2、ViT-B/16、DeiT-Tiny 和 Swin-Tiny 等模型，并与随机剪枝、量级剪枝、静态显著性剪枝以及预算贪婪选择进行了对比。
Across five seeds, the bandit methods usually reduce degradation relative to budgeted greedy in the paired language-model comparisons. Of 28 comparisons highlighted in the paper, 23 bootstrap confidence intervals exclude zero and 11 paired tests have p < 0.05; six have q < 0.05 after Benjamini-Hochberg correction across the full family of 116 dataset-wise tests.
在五个随机种子的实验中，在配对语言模型比较中，老虎机方法通常比预算贪婪法更能减少性能下降。在论文强调的 28 组比较中，23 个自助法（bootstrap）置信区间不包含零，11 个配对检验的 p 值小于 0.05；在针对全部 116 个数据集测试进行 Benjamini-Hochberg 校正后，有 6 个测试的 q 值小于 0.05。
Matched-evaluation results for ViT-B/16 and Swin-Tiny indicate that their gains are not explained solely by a larger candidate-evaluation budget.
针对 ViT-B/16 和 Swin-Tiny 的匹配评估结果表明，它们的性能提升并不能仅仅归因于更大的候选评估预算。