I am currently experimenting with weight pruning to optimize my Transformer models for edge deployment. While I’ve managed to reduce the model size by nearly 40%, I am noticing a slight drop in the BLEU scores for my translation tasks. How does weight pruning specifically affect the performance of large-scale transformers, and is there a sweet spot for sparsity without losing precision?
3 answers
Weight pruning in Transformer models primarily functions by removing redundant parameters, which can lead to faster inference and lower memory usage. In my experience with BERT and GPT-based architectures, unstructured pruning often causes a more noticeable drop in accuracy compared to structured pruning because it disrupts the hardware-level optimizations. However, if you implement "Lottery Ticket Hypothesis" principles, you can find sub-networks that maintain nearly 98% of the original performance. The "sweet spot" usually lies around 30% to 50% sparsity; beyond that, the attention mechanisms start to lose the ability to capture complex semantic relationships, leading to the score drops you are seeing.
That is a great observation regarding the BLEU scores. Have you considered whether you are using iterative magnitude pruning or a one-shot approach during your fine-tuning phase? Also, are you applying a global pruning mask or layer-wise?
In my projects, weight pruning usually works best when followed by a brief period of retraining or "fine-tuning" to allow the remaining weights to compensate for the lost connections.
I totally agree with this. Fine-tuning after pruning is non-negotiable for Transformers. I would also suggest looking into knowledge distillation as a complementary technique to recover those lost BLEU points.
I’ve been using iterative magnitude pruning across the entire model globally. I noticed that the encoder layers seem more resilient than the decoder layers. Does the specific layer type usually dictate how much sparsity we can actually get away with before the model’s linguistic understanding begins to collapse?