Pruning Attention Heads of Transformer Models Using A* Search: A Novel Approach to Compress Big NLP Architectures

TL;DR

This paper introduces a novel A* search-based pruning method to eliminate redundant attention heads in Transformer models like BERT, significantly reducing model size without sacrificing accuracy.

Contribution

It presents a new pruning algorithm using A* search that guarantees accuracy while removing up to 40% of attention heads in Transformer models.

Findings

Up to 40% attention head removal with no accuracy loss

A* search guarantees pruning with accuracy preservation

Effective compression method for large NLP models

Abstract

Recent years have seen a growing adoption of Transformer models such as BERT in Natural Language Processing and even in Computer Vision. However, due to their size, there has been limited adoption of such models within resource-constrained computing environments. This paper proposes novel pruning algorithm to compress transformer models by eliminating redundant Attention Heads. We apply the A* search algorithm to obtain a pruned model with strict accuracy guarantees. Our results indicate that the method could eliminate as much as 40% of the attention heads in the BERT transformer model with no loss in accuracy.