[Paper Review] DynaBERT: Dynamic BERT with Adaptive Width and Depth

Lu Hou, Zhiqi Huang+4 more|arXiv (Cornell University)|Apr 8, 2020

Topic Modeling40 references119 citations

TL;DR

DynaBERT trains a width- and depth-adaptive BERT model that can deploy sub-networks with varying resource usage, using knowledge distillation and network rewiring to maintain accuracy.”

ABSTRACT

The pre-trained language models like BERT, though powerful in many natural language processing tasks, are both computation and memory expensive. To alleviate this problem, one approach is to compress them for specific tasks before deployment. However, recent works on BERT compression usually compress the large BERT model to a fixed smaller size. They can not fully satisfy the requirements of different edge devices with various hardware performances. In this paper, we propose a novel dynamic BERT model (abbreviated as DynaBERT), which can flexibly adjust the size and latency by selecting adaptive width and depth. The training process of DynaBERT includes first training a width-adaptive BERT and then allowing both adaptive width and depth, by distilling knowledge from the full-sized model to small sub-networks. Network rewiring is also used to keep the more important attention heads and neurons shared by more sub-networks. Comprehensive experiments under various efficiency constraints demonstrate that our proposed dynamic BERT (or RoBERTa) at its largest size has comparable performance as BERT-base (or RoBERTa-base), while at smaller widths and depths consistently outperforms existing BERT compression methods. Code is available at https://github.com/huawei-noah/Pretrained-Language-Model/tree/master/DynaBERT.

Motivation & Objective

• Motivate efficient deployment of BERT on edge devices with heterogeneous hardware.
• Enable flexible inference by selecting sub-networks of varying width and depth within one model.
• Maintain high accuracy across configurations through knowledge distillation from a full-size teacher.
• Bridge size gaps between teacher and student with a width-adaptive teacher-assisted approach.
• Demonstrate that dynamic sub-networks can match or exceed fixed-size compression baselines on standard benchmarks.

Proposed method

• Define width and depth adaptivity in Transformer layers by varying the number of attention heads and FFN neurons (width) and the number of Transformer layers (depth).
• Two-stage training: first train a width-adaptive DynaBERT_W with network rewiring and distillation from a rewired teacher; then train width-and-depth-adaptive DynaBERT using DynaBERT_W as teacher.
• Use network rewiring to rank and align important heads/neurons so they are shared across sub-networks.
• Apply knowledge distillation across logits, embeddings, and hidden states to transfer knowledge from teacher to student across widths and depths.
• For DynaBERT_W, use a fixed rewired teacher and multiple width configurations; for full DynaBERT, distill from the widest-depth teacher to sub-networks with reduced depth and width.
• Data augmentation from TinyBERT to provide task-specific examples for distillation.

Experimental results

Research questions

• RQ1 Can a single BERT model be trained to support multiple width and depth configurations with minimal accuracy loss compared to fixed-size models?
• RQ2 Does expanding training with width and depth adaptivity yield better efficiency-accuracy trade-offs than fixed-size compression methods?
• RQ3 What is the impact of network rewiring and distillation on performance across GLUE and SQuAD tasks?
• RQ4 Is using a width-adaptive model as a teacher assistant beneficial for bridging the gap between large and small sub-networks?
• RQ5 How do DynaBERT_W and DynaBERT compare against state-of-the-art fixed-size distillation and depth-only adaptive methods under various efficiency constraints?

Key findings

• DynaBERT at its largest size achieves comparable performance to BERT_BASE on benchmarks when configured for full width/depth.
• Sub-networks with reduced width and/or depth consistently outperform fixed-size compression methods under identical efficiency constraints.
• Width adaptivity is more robust than depth alone, and DynaBERT_W (width-adaptive) generally outperforms DynaBERT_D (depth-adaptive) for smaller configurations.
• DynaBERT sub-networks often surpass DistilBERT, TinyBERT, and LayerDrop under similar resource budgets, across GLUE and SQuAD tasks.
• Using DynaBERT_W as a teacher assistant improves distillation performance over directly distilling from the rewired full model.
• No further fine-tuning is required for each sub-network after training, enabling efficient deployment across devices.