Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

VTrans: Accelerating Transformer Compression with Variational Information Bottleneck based Pruning

Oshin Dutta, Ritvik Gupta, Sumeet Agarwal

TL;DR

VTrans introduces a unified, VIB-guided pruning framework that compresses embeddings, attention heads, and layers in transformers to meet user-defined parameter or FLOPs constraints. It employs stochastic VIB masks and distillation to prune redundancies while preserving predictive performance, and adds Faster-VTrans variants that use data subsets or mask-only training to dramatically accelerate compression. Across BERT, RoBERTa, GPT-2, and LLaMA-2-7B, it achieves up to ~70–80% compression with minimal accuracy loss and substantial speedups, while providing qualitative evidence of reduced redundancy in attention heads. The approach demonstrates strong scalability and interpretability, making it practical for deploying compact, efficient transformers on edge devices and large-scale models alike.

Abstract

In recent years, there has been a growing emphasis on compressing large pre-trained transformer models for resource-constrained devices. However, traditional pruning methods often leave the embedding layer untouched, leading to model over-parameterization. Additionally, they require extensive compression time with large datasets to maintain performance in pruned models. To address these challenges, we propose VTrans, an iterative pruning framework guided by the Variational Information Bottleneck (VIB) principle. Our method compresses all structural components, including embeddings, attention heads, and layers using VIB-trained masks. This approach retains only essential weights in each layer, ensuring compliance with specified model size or computational constraints. Notably, our method achieves upto 70% more compression than prior state-of-the-art approaches, both task-agnostic and task-specific. We further propose faster variants of our method: Fast-VTrans utilizing only 3% of the data and Faster-VTrans, a time efficient alternative that involves exclusive finetuning of VIB masks, accelerating compression by upto 25 times with minimal performance loss compared to previous methods. Extensive experiments on BERT, ROBERTa, and GPT-2 models substantiate the efficacy of our method. Moreover, our method demonstrates scalability in compressing large models such as LLaMA-2-7B, achieving superior performance compared to previous pruning methods. Additionally, we use attention-based probing to qualitatively assess model redundancy and interpret the efficiency of our approach. Notably, our method considers heads with high attention to special and current tokens in un-pruned model as foremost candidates for pruning while retained heads are observed to attend more to task-critical keywords.

VTrans: Accelerating Transformer Compression with Variational Information Bottleneck based Pruning

TL;DR

VTrans introduces a unified, VIB-guided pruning framework that compresses embeddings, attention heads, and layers in transformers to meet user-defined parameter or FLOPs constraints. It employs stochastic VIB masks and distillation to prune redundancies while preserving predictive performance, and adds Faster-VTrans variants that use data subsets or mask-only training to dramatically accelerate compression. Across BERT, RoBERTa, GPT-2, and LLaMA-2-7B, it achieves up to ~70–80% compression with minimal accuracy loss and substantial speedups, while providing qualitative evidence of reduced redundancy in attention heads. The approach demonstrates strong scalability and interpretability, making it practical for deploying compact, efficient transformers on edge devices and large-scale models alike.

Abstract

In recent years, there has been a growing emphasis on compressing large pre-trained transformer models for resource-constrained devices. However, traditional pruning methods often leave the embedding layer untouched, leading to model over-parameterization. Additionally, they require extensive compression time with large datasets to maintain performance in pruned models. To address these challenges, we propose VTrans, an iterative pruning framework guided by the Variational Information Bottleneck (VIB) principle. Our method compresses all structural components, including embeddings, attention heads, and layers using VIB-trained masks. This approach retains only essential weights in each layer, ensuring compliance with specified model size or computational constraints. Notably, our method achieves upto 70% more compression than prior state-of-the-art approaches, both task-agnostic and task-specific. We further propose faster variants of our method: Fast-VTrans utilizing only 3% of the data and Faster-VTrans, a time efficient alternative that involves exclusive finetuning of VIB masks, accelerating compression by upto 25 times with minimal performance loss compared to previous methods. Extensive experiments on BERT, ROBERTa, and GPT-2 models substantiate the efficacy of our method. Moreover, our method demonstrates scalability in compressing large models such as LLaMA-2-7B, achieving superior performance compared to previous pruning methods. Additionally, we use attention-based probing to qualitatively assess model redundancy and interpret the efficiency of our approach. Notably, our method considers heads with high attention to special and current tokens in un-pruned model as foremost candidates for pruning while retained heads are observed to attend more to task-critical keywords.
Paper Structure (41 sections, 4 equations, 13 figures, 15 tables, 1 algorithm)

This paper contains 41 sections, 4 equations, 13 figures, 15 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Related Work
  3. Background
  4. Basic Transformer Architecture
  5. Compressed Representations with Variational Information Bottleneck
  6. Method
  7. VIB-based Pruning
  8. Knowledge Distillation
  9. Faster variants of our method
  10. Experiments
  11. Implementation Details
  12. Evaluation
  13. Feasibility of Compression with Faster Variants
  14. Scaling-up to prune more than billion parameters
  15. Ablation Study
  16. ...and 26 more sections

Figures (13)

  • Figure 1: (a) Our primary method - VTrans involves training all pre-trained model parameters and VIB masks on the entire dataset during pruning, but during finetuning, only the unmasked important weights are updated. (b) Fast-VTrans utilizes a subset of data for both pruning and finetuning. (c) Faster-VTrans, the fastest among all, trains VIB masks, Add Norm layer, and model bias parameters during pruning and finetuning, using only a subset of the entire dataset. (d) The masked and redundant units are removed after the finetuning phase, resulting in the dense compressed model.
  • Figure 2: Our method surpasses previous techniques in compressing BERT-base, ROBERTa-base, and GPT-2-small models. Our faster variants have reduced compression timeframes. All models have 28M parameters being compressed from their respective teachers.
  • Figure 3: Pruning BERT-base (teacher model) on GLUE and SQuAD tasks: Each point denotes the mean performance of pruned models averaged over five random trials with different seeds with average variance of $\pm 0.4\%$ across all models. The dashed line represents a 1% performance drop from the teacher model.
  • Figure 4: Speeding-up Compression: Faster-VTrans achieves significant speed-up over VTrans on GLUE tasks, with minimal performance loss for models 35% smaller than GPT-2$_{small}$.
  • Figure 5: (a) Comparative attention allocation towards special tokens in un-pruned (top) and pruned (bottom) models. Dots represent token attention by heads, with lines indicating mean attention by remaining heads. Pruned model exhibits reduced special token attention (b) Eliminated heads example: those highly attentive to CLS, SEP, current token, common articles like 'the' and displaying broad attention. Line thickness reflects attention weight towards the token.
  • ...and 8 more figures