PaperlayerView exact revision on arXiv

Source-linked AI summary

Damage-Aware Bandit Pruning for Vision and Language Transformers

Salem Ameen, Sunil Vadera

arXiv:2609.05448v1cs.AIcs.LG

TL;DR

Transformer pruning needs budget-efficient ways to identify complete units whose removal causes limited degradation. This paper uses paired damage-aware bandit selection with sequential masking for attention heads and MLP groups, finding gains in several language and matched-evaluation vision configurations while evaluating functional suppression rather than realised compression.

  • Problem

    Transformer pruning requires selecting removable complete units without exhaustively testing every candidate on large calibration sets.

  • Method

    The method temporarily masks attention heads or MLP groups, converts paired same-batch loss damage into bounded rewards, and sequentially constructs a top-k mask using MAB-UCB or MAB-TS.

  • Results

    Bandit selection frequently reduces degradation relative to budgeted greedy, including matched-evaluation gains for ViT-B/16 and Swin-Tiny; 6 of 116 corrected dataset-wise tests have q < 0.05.

  • Takeaways & Limitations

    The findings support improved structured-unit selection in several configurations, but not a universal advantage across all settings.

  • Takeaways & Limitations

    Functional zeroing leaves dense checkpoint tensors stored, so reported parameter effects do not measure physical compression, memory reduction, or deployment speedup.

Abstract

from arXiv · show

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. 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. 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. 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. Matched-evaluation results for ViT-B/16 and Swin-Tiny indicate that their gains are not explained solely by a larger candidate-evaluation budget.

1. Introduction

The paper frames structured transformer-unit pruning as fixed-budget bandit selection, using paired damage and bounded rewards to choose attention heads and MLP groups. It evaluates this approach across language and vision models against several baselines while distinguishing functional suppression from physical compression.

  • Structured pruning selects complete attention heads and MLP groups as candidate computational units.The motivation is that many transformer heads can be removed with limited degradation, but exhaustive candidate testing is impractical.
  • The method treats unit selection as a fixed-budget multi-armed bandit problem with sequential top-k mask construction.MAB-UCB and fractional-Beta MAB-TS allocate repeated evaluations among candidates while conditioning on units already selected.
  • Selected units are functionally zeroed in the original dense checkpoint, so the experiments assess effective suppression and selection quality rather than physical compression or speedup.Dense tensors remain stored.
  • Paired same-batch damage compares the temporarily masked model with the base or current masked model, reducing batch-to-batch variation in unit comparisons.The resulting bounded reward drives both bandit policies.
  • The study compares bandit selection with random, magnitude, Wanda-style, Taylor, Fisher, and budgeted-greedy baselines across five seeds.Experiments cover WikiText-2, LAMBADA, and Imagenette, including GPT-style language models and ViT, DeiT, and Swin transformers.

2. Related work and positioning

The paper positions its contribution as budget-aware selection of structured transformer units rather than physical architecture reduction. It extends earlier bandit pruning work with typed transformer candidates, paired damage, and sequential conditioning.

  • Earlier studies established uneven head importance and removable transformer-head redundancy but did not allocate a limited evaluation budget under noisy candidate effects.The paper addresses this selection problem with adaptive bandit search.
  • Existing pruning methods target one-shot weight sparsity, structured masks under FLOP or latency constraints, gradient-based coupled structures, activation fluctuations, or dense matrix dimension reduction.These methods pursue objectives or structural changes distinct from the present budget-allocation focus.
  • The authors' earlier work applied fixed-budget MAB selection to scalar weights, structured neurons, and convolutional feature maps, providing the methodological foundation here.The present study transfers that framework to transformer structured units.
  • The present method uses attention heads or contiguous MLP groups, paired same-batch damage under the cumulative mask, and an explicit run-specific evaluation budget.Unlike methods that physically reduce structures, it retains the dense checkpoint and uses functional zeroing.

3. Method

The method treats structured transformer-unit selection as a fixed-budget, sequential bandit problem using paired damage and bounded rewards. Attention heads and MLP groups are functionally masked in the dense model, so parameter effects represent effective suppression rather than physical compression.

  • Candidate construction: Attention heads and contiguous MLP channel groups form typed candidate units for structured transformer pruning.
  • Sequential selection: At each step, previously selected units remain suppressed while a new candidate is temporarily masked and evaluated on the same calibration batch.This produces marginal damage conditioned on the current selected set rather than jointly searching all combinations.
  • Parameter accounting: Functional zeroing suppresses selected forward contributions while retaining dense checkpoint tensors, so reported reductions are effective zeroed-parameter estimates.They do not measure physical checkpoint-size reduction or deployment speed-up.
  • Damage and reward: Paired damage compares masked and current-base losses on identical examples, reducing batch variation and distinguishing harmful, neutral, and beneficial removals.
  • Bandit policies: A bounded reward drives MAB-UCB or fractional-Beta Thompson Sampling, with UCB using empirical reward and exploration statistics over a fixed pull budget.Unevaluated active candidates are sampled first when the pull budget permits.
  • Mask construction: The final mask is constructed sequentially by adding one unit at a time until k units are selected.The procedure does not delete several candidates simultaneously.

4. Experimental setup

The experiments compare damage-aware bandit selection with alternative pruning techniques across language and vision transformers, datasets, and metrics under controlled evaluation protocols. Results are reported across seeds where stochastic selection or calibration sampling varies, with explicit accounting for evaluation budgets and effective functional zeroing.

  • Reproducibility: Random, greedy, and bandit methods are reported over seeds when their selection process or sampled calibration batches vary across runs.Unless otherwise stated, stochastic results use five seeds: 1, 2, 3, 42, and 123.
  • Compared methods: MAB-UCB and MAB-TS use paired trials with damage-aware UCB or Thompson-style selection, respectively.The compared techniques distinguish calibration use from counted paired candidate-evaluation trials.
  • Models and datasets: The study covers GPT-2, OPT, Pythia, Qwen2.5-0.5B, SmolLM2-360M, ViT-B/16, DeiT-Tiny, and Swin-Tiny on WikiText-2, LAMBADA, and Imagenette.WikiText-2 supports main language evaluation and calibration, LAMBADA provides cross-dataset evaluation, and Imagenette supports vision-transformer experiments.
  • Evaluation protocol: The main language configuration uses sequence length 128, up to 512 calibration texts, up to 1024 evaluation texts, and 80 evaluation batches.Vision runs use Imagenette image-folder data, batch size 8, up to 1024 calibration images, and up to 2000 evaluation images.
  • Evaluation budget: The language setup uses a maximum pre-screened pool of 48 candidates, 32 paired evaluations per step, two calibration mini-batches per damage estimate, and 32 greedy trials.The configuration also sets c_ucb = 1.5, MLP group size g_0 = 32, and reward temperature τ = 0.02.
  • Vision accounting: Vision-transformer runs report each method’s total paired candidate-evaluation budget explicitly, while nominal pruning ratios are converted to integer head counts.For DeiT-Tiny, a nominal 0.10 ratio selects 4 of 36 heads, or 11.11%; Swin-Tiny reports effective zeroed-parameter percentages as a range.

5. Results

Across language and vision experiments, damage-aware bandit selection often preserved performance better than budgeted greedy under comparable pruning settings. The evidence includes strong GPT-2 and LAMBADA gains, matched-budget vision improvements, and statistically supportive but non-universal language-model comparisons.

  • MAB methods often kept perplexity degradation low, while several static saliency methods degraded substantially on particular model families.
  • GPT-2 MLP-group pruning: At 5% GPT-2 MLP pruning, MAB-UCB produced -4.85% percent perplexity change versus -1.99% for budgeted greedy; at 8%, it produced -1.65% versus +2.05%.MAB-UCB was best at 3%, 5%, and 8% tested ratios, with negative perplexity changes.
  • LAMBADA cross-dataset evaluation: On LAMBADA, MAB gave the lowest degradation for all four models; Pythia-410M reached 30.30% perplexity increase versus 238.85% for budgeted greedy.For Qwen2.5-0.5B, MAB-UCB reached 10.46%, versus 30.88% for magnitude and 49.90% for budgeted greedy.
  • Matched-budget vision-transformer evaluation: Under matched paired-evaluation budgets, MAB-TS reduced the Top-1 drop from 3.02 to 2.05 percentage points on ViT-B/16 and from 2.43 to 1.34 on Swin-Tiny.On Swin-Tiny, MAB-UCB yielded the lowest reported cross-entropy loss; DeiT-Tiny was not used as the primary fairness claim because its greedy run used fewer trials.
  • Scope of reported effects: The experiments measure effective structured suppression through functional zeroing in dense checkpoints, not physical compression, memory reduction, or deployment speedup.The vision tables report dense and zeroed settings, while Figure 4 omits magnitude because severe Swin-Tiny collapse dominates the scale.
  • Statistical analysis and reproducibility: Among 28 highlighted MAB-versus-greedy language comparisons, 23 bootstrap confidence intervals excluded zero, 11 paired tests had p < 0.05, and 6 had q < 0.05 after correction across 116 tests.The paper states these results support improved selection in several configurations rather than a universal advantage; non-finite SmolLM2 runs indicate model- and setting-dependent stability.

6. Discussion

The bandit methods often improve structured-unit selection under bounded evaluation budgets, but their advantage is configuration-dependent and deployment value remains unresolved.

  • Empirical findings: 11 of 28 highlighted language comparisons have paired tests with p < 0.05, while 6 remain below q < 0.05 after correction across 116 tests.The evidence supports improvements in several configurations rather than a universal advantage.
  • Empirical findings: For ViT-B/16 and Swin-Tiny, bandit selection yields lower reported degradation when greedy and MAB use the same trial count.DeiT-Tiny is excluded from this fairness claim because its greedy run uses fewer evaluations.
  • Interpretation: Adaptive allocation can be useful under a bounded search budget, although the limited benchmark scale and five-seed design constrain generalisation.The conclusion is supported jointly by language and vision experiments.
  • Positioning: The study prioritizes evaluation-budget allocation and paired damage over physical architecture reduction, extending earlier MAB pruning with typed units and sequential mask conditioning.Its deployment value still requires physical model surgery, broader baselines, larger benchmarks, and hardware measurements.

7. Limitations

The study’s evidence is constrained by functional masking, limited visual and language coverage, stability failures, and search assumptions that leave important pruning effects outside the optimization.

  • Deployment scope: Functional zeroing leaves dense checkpoint tensors stored, so effective zeroed-parameter counts do not establish memory reduction, speedup, or hardware acceleration.Those deployment benefits require separate model surgery and hardware evaluation.
  • Benchmark scope: Imagenette, rather than full ImageNet-1K, limits the scale of the vision-transformer evidence.The protocol should be repeated on ImageNet-1K or another large-scale visual-recognition benchmark.
  • Stability: SmolLM2 has non-finite perplexity in some configurations, which are reported as stability failures rather than averaged away.This indicates that stability depends on the model and pruning setting.
  • Search assumptions: The reward estimates marginal single-unit damage, so interactions among multiple selected units are not explicitly optimized during search.Such interactions are assessed only by the final empirical evaluation.
  • Search assumptions: The method covers post-training pruning without fine-tuning or recovery training and restricts the main language search to a low-magnitude pre-screened pool of at most 48 candidates.These choices constrain the method’s evaluated scope.

8. Conclusion

The paper frames transformer structured-unit selection as a fixed-budget damage-aware bandit problem and reports reduced degradation in several language and vision configurations.

  • Conclusion: The method evaluates attention heads and MLP groups using paired same-batch loss changes, then constructs masks sequentially with MAB-UCB or MAB-TS.The paired design measures masked damage relative to the current masked model.
  • Conclusion: Bandit selection reduces degradation in several reported experiments, including matched-evaluation results for ViT-B/16 and Swin-Tiny.Only a subset of highlighted paired comparisons remains significant after multiple-testing correction.
  • Conclusion: Functional zeroing means the findings concern selection quality rather than realised compression or speedup.Dense checkpoint tensors remain present during the experiments.
  • Conclusion: Future work should add physical model surgery, latency and memory measurements, broader structured-pruning comparisons, ImageNet-1K, and larger language models.These directions address the paper’s stated deployment and generalisation boundaries.

Appendix A. Complete detailed result tables

The appendix tables provide complete per-method records underlying the paper’s summary results.

  • Appendix: Appendix Tables A1-A5 report final performance, percentage change, pruning target and ratio, effective zeroed-parameter estimates, and evaluation counts.The records include final perplexity or accuracy as applicable.

Appendix Table A1. Detailed WikiText-2 10% head-pruning results.

Appendix Table A1 reports WikiText-2 10% attention-head pruning outcomes across multiple language models and selection methods, including final perplexity, perplexity change, and effective suppression.

  • GPT-2: 13.35% PPL change is reported for GPT-2 heads selected by Taylor, compared with 14.52% for Fisher and 33.91% for Wanda-style selection.The corresponding final perplexities are 71.4956, 72.2332, and 84.4646, respectively.
  • GPT-2-medium: 8.38% PPL change is reported for GPT-2-medium heads selected by Taylor, compared with 8.78% for Fisher.Their final perplexities are 48.7872 and 48.9653, respectively.
  • OPT: 5.78% PPL change is reported for OPT-125m heads selected by Fisher, compared with 6.05% for Taylor and 7.52% for magnitude selection.The final perplexity for Fisher is 78.5348.
  • OPT: 7.96% PPL change is reported for OPT-350m heads selected by Fisher, compared with 12.60% for Taylor and 36.75% for Wanda-style selection.The final perplexity for Fisher is 63.5288.
  • Pythia: 23.27% PPL change is reported for Pythia-160m heads selected by Taylor, compared with 70.95% for Fisher and 106.44% for magnitude selection.The corresponding final perplexities are 85.0937, 118.0075, and 142.4696.
  • Qwen2.5: 35.36% PPL change is reported for Qwen2.5-0.5B heads selected by magnitude, compared with 66.08 ± 53.19% for random selection.Wanda-style, Taylor, and Fisher selection report larger PPL changes of 21298.44%, 302.34%, and 356.85%, respectively.

Appendix Table B1. Representative selection-cost and fairness summary for All-LLM 10% head-pruning rows.

Appendix Table B1 separates calibration use from counted paired candidate-evaluation trials, while Table B2 summarizes statistical evidence for highlighted MAB-versus-greedy language-model comparisons.

  • Selection-cost accounting: Random and magnitude selection do not perform candidate testing, whereas budgeted greedy and MAB methods use the same counted candidate-evaluation convention.Wanda-style, Taylor, and Fisher use calibration information but are not adaptive paired candidate-evaluation procedures.
  • Statistical summary: 28 highlighted paired language-model comparisons against budgeted greedy are summarized in Appendix Table B2.The summary uses the full 116-test family for Benjamini-Hochberg q-values.
  • Statistical summary: 23/28 bootstrap 95% confidence intervals exclude zero.This is the largest reported count among the three statistical criteria summarized for the highlighted comparisons.
  • Statistical summary: 11/28 paired tests have p < 0.05, while 6/28 have BH q < 0.05 across the 116-test family.The corrected count is smaller than the nominal significance count.
  • Imagenette comparison: Lower Top-1 drop is better in the full Imagenette comparison after functional head zeroing.The figure includes Random, Magnitude, Budgeted greedy, MAB-UCB, and MAB-TS.
Loading 2609.05448v1…