BHPVAS: visual analysis system for pruning attention heads in BERT model

Abstract

In the field of deep learning, pre-trained BERT models have achieved remarkable success. However, the accompanying problem is that models with more complex structures and more network parameters. The huge parameter size makes the computational cost in terms of time and memory become extremely expensive. Recent work has indicated that BERT models own a significant amount of redundant attention heads. Meanwhile considerable BERT models compression algorithms have been proposed, which can effectively reduce model complexity and redundancy with pruning some attention heads. Nevertheless, existing automated model compression solutions are mainly based on predetermined pruning program, which requires multiple expensive pruning-retraining cycles or heuristic designs to select additional hyperparameters. Furthermore, the training process of BERT models is a black box, and lacks interpretability, which makes researchers cannot intuitively understand the optimization process of the model. In this paper, we propose a visual analysis system, BHPVAS, for pruning BERT models, which helps researchers to incorporate their understanding of model structure and operating mechanism into the model pruning process and generate pruning schemes. We propose three pruning criteria based on the attention data, namely, importance score, stability score, and similarity score, for evaluating the importance of self-attention heads. Additionally, we design multiple collaborative views to display the entire pruning process, guiding users to carry out pruning. Our system supports exploring the role of self-attention heads in the model inference process using text dependency relations and attention weight distribution. Finally, we conduct two case studies to demonstrate how to use the system for Sentiment Classification Sample Analysis and Pruning Scheme Exploration, verifying the effectiveness of the visual analysis system.

Graphical Abstract