Damage-Aware Bandit Pruning for Vision and Language Transformers

Summary

The paper formulates structured post-training pruning for language and vision Transformers as a damage-aware multi-armed bandit problem. Attention heads and MLP channel groups are temporarily masked on calibration batches, and paired damage is measured as the masked loss minus the base loss on the same batch to reduce batch-to-batch variation. A smooth bounded reward drives either a UCB-style policy or fractional-Beta Thompson Sampling, while the final mask is built sequentially by adding one unit at a time. The selected units are zeroed in the original dense checkpoint, so the reported parameter effects describe effective structural suppression rather than physical compression or measured inference speedup. Experiments use WikiText-2, LAMBADA, and Imagenette with GPT-2, OPT, Pythia, Qwen2.5, SmolLM2, ViT-B/16, DeiT-Tiny, and Swin-Tiny, comparing the bandit methods with random, magnitude, static-saliency, and budgeted-greedy selection. Across five seeds, bandit methods usually reduce degradation relative to budgeted greedy in paired language-model comparisons. Among 28 highlighted comparisons, 23 bootstrap confidence intervals exclude zero, 11 paired tests have p < 0.05, and six remain significant at q < 0.05 after Benjamini-Hochberg correction across 116 dataset-wise tests. Matched-budget results for ViT-B/16 and Swin-Tiny suggest the vision gains are not explained only by a larger candidate-evaluation budget.