[Paper Review] Structured Pruning of a BERT-based Question Answering Model

Jason S. McCarley, Rishav Chakravarti+1 more|arXiv (Cornell University)|Oct 14, 2019

Topic Modeling36 references72 citations

TL;DR

The paper shows that task-specific structured pruning combined with distillation can compress BERT/RoBERTa QA models, yielding substantial speedups with minimal accuracy loss without re-pretraining.

ABSTRACT

The recent trend in industry-setting Natural Language Processing (NLP) research has been to operate large %scale pretrained language models like BERT under strict computational limits. While most model compression work has focused on "distilling" a general-purpose language representation using expensive pretraining distillation, less attention has been paid to creating smaller task-specific language representations which, arguably, are more useful in an industry setting. In this paper, we investigate compressing BERT- and RoBERTa-based question answering systems by structured pruning of parameters from the underlying transformer model. We find that an inexpensive combination of task-specific structured pruning and task-specific distillation, without the expense of pretraining distillation, yields highly-performing models across a range of speed/accuracy tradeoff operating points. We start from existing full-size models trained for SQuAD 2.0 or Natural Questions and introduce gates that allow selected parts of transformers to be individually eliminated. Specifically, we investigate (1) structured pruning to reduce the number of parameters in each transformer layer, (2) applicability to both BERT- and RoBERTa-based models, (3) applicability to both SQuAD 2.0 and Natural Questions, and (4) combining structured pruning with distillation. We achieve a near-doubling of inference speed with less than a 0.5 F1-point loss in short answer accuracy on Natural Questions.

Motivation & Objective

• Motivate reducing computational burden of large QA models in industry settings without retraining from scratch.
• Develop and evaluate task-specific structured pruning for BERT/RoBERTa QA systems.
• Investigate combining structured pruning with distillation to maximize speedups with minimal accuracy loss.
• Assess transferability of pruning techniques across datasets (SQuAD 2.0 and Natural Questions) and models (BERT and RoBERTa).
• Provide guidelines on gate placement and pruning strategies for transformer-based QA models.

Proposed method

• Insert trainable gate masks into each transformer layer to selectively deactivate attention heads and feed-forward activations.
• Explore gate-determination strategies including random, head importance gain, and L0 regularization with hard-concrete distributions.
• Prune attention heads and feed-forward components by applying gates and retraining the remaining parameters with or without distillation.
• Continue task-specific training after pruning, with options for standard retraining or distillation from an unpruned teacher model.
• Apply pruning to both BERT-based and RoBERTa-based QA systems and evaluate on SQuAD 2.0 and Natural Questions.
• Demonstrate that non-identical, layer-wise pruning structures can emerge (lighter near the ends, more complex in the middle).

Experimental results

Research questions

• RQ1 Can task-specific structured pruning transfer from BERT-base to BERT-large in QA tasks?
• RQ2 Do structured pruning techniques developed for BERT-based QA transfer to RoBERTa and to Natural Questions?
• RQ3 Does combining structured pruning with distillation yield better speed/accuracy tradeoffs than pruning or distillation alone?
• RQ4 What are effective gate-determination strategies (random, gain, L0) for pruning QA transformers?
• RQ5 How does pruning affect long-answer vs short-answer QA performance on SQuAD 2.0 and Natural Questions?

Key findings

• L0 regularization pruning enables substantial pruning (up to ~48% attention heads and ~70% feed-forward activations) with less than 5 F1 points loss on SQuAD 2.0.
• Combined pruning and continued training recovers much of the accuracy and nearly doubles decoding speed on SQuAD 2.0.
• Pruned RoBERTa-large on Natural Questions achieves notable speedups with small F1 loss and benefits from distillation.
• Distillation after pruning yields additional gains, approaching the performance of larger baselines while maintaining faster inference.
• Pruned models tend to have non-identical layer structures, with intermediate layers retained more heavily than the extremes, suggesting a non-uniform pruning pattern across layers.
• Pruning techniques transfer to RoBERTa and to NQ, indicating model- and dataset-agnostic robustness of the approach.