AUBER: Automated BERT regularization

Abstract

How can we effectively regularize BERT? Although BERT proves its effectiveness in various NLP tasks, it often overfits when there are only a small number of training instances. A promising direction to regularize BERT is based on pruning its attention heads with a proxy score for head importance. However, these methods are usually suboptimal since they resort to arbitrarily determined numbers of attention heads to be pruned and do not directly aim for the performance enhancement. In order to overcome such a limitation, we propose AUBER, an automated BERT regularization method, that leverages reinforcement learning to automatically prune the proper attention heads from BERT. We also minimize the model complexity and the action search space by proposing a low-dimensional state representation and dually-greedy approach for training. Experimental results show that AUBER outperforms existing pruning methods by achieving up to 9.58% better performance. In addition, the ablation study demonstrates the effectiveness of design choices for AUBER.

Fig 1. Performance of AUBER and its competitors on 4 GLUE datasets.

AUBER successfully regularizes BERT model, enhancing the model performance up to 9.58%. AUBER provides the best performance among the state-of-the-art BERT attention head pruning methods.

Fig 2. Overview of transitioning in AUBER.

The figure shows the transition from Layer 2 to Layer 3 in AUBER with BERT-base.

Fig 3. Overall flow of training AUBER on a layer.

AUBER trains DQN to find out the attention heads that should be pruned for a better regularization following the illustrated steps.

Fig 4. Performance after pruning the attention heads from each layer.

AUBER consistently improves the model performance and achieves outstanding final performance, while all the other methods fail to enhance the model performance.