Structured post‑training pruning must select complete functional units whose suppression causes limited performance loss. We cast the selection of units in language and vision Transformers as a damage‑aware multi‑armed bandit problem under a fixed candidate‑evaluation budget.
Attention heads and MLP channel groups are temporarily masked on calibration batches. The paired damage is defined as the masked loss minus the base loss on the same batch, which reduces batch‑to‑batch variance.
A smooth bounded reward drives either a UCB‑style policy or fractional‑Beta Thompson Sampling, and the final mask is built sequentially by adding one unit at each step. The selected units are functionally zeroed in the original dense checkpoint, so the reported parameter effects represent effective structural suppression rather than physical compression or measured speedup.
Experiments on WikiText‑2, LAMBADA and Imagenette involve GPT‑2, OPT, Pythia, Qwen2.5, SmolLM2, ViT‑B/16, DeiT‑Tiny and Swin‑Tiny, with baselines including random, magnitude, static‑saliency and budgeted‑greedy selection. Across five seeds, the bandit methods usually reduce degradation relative to budgeted greedy in paired language‑model comparisons. Of the 28 highlighted comparisons, 23 bootstrap confidence intervals exclude zero and 11 paired tests achieve significant p‑values.
Review