PaperlayerView exact revision on arXiv

Source-linked AI summary

DynaBERT: Dynamic BERT with Adaptive Width and Depth

Lu Hou, Zhiqi Huang, Lifeng Shang, Xin Jiang, Xiao Chen, Qun Liu

arXiv:2004.04037v2cs.CLcs.LG

TL;DR

BERT’s high resource cost and fixed-size compression limit deployment across heterogeneous and changing edge-device conditions. DynaBERT distills knowledge through staged training of width- and depth-adaptive sub-networks, with rewiring to share important components. Across GLUE and SQuAD, its sub-networks match or exceed baseline performance under comparable efficiency constraints.

  • Problem

    BERT is difficult to deploy on edge devices because its resource demands are high, while fixed-size compression cannot accommodate varied hardware and changing resource conditions.

  • Method

    DynaBERT trains a width-adaptive model and then a width- and depth-adaptive model through knowledge distillation, using rewiring to share important attention heads and neurons.

  • Results

    Under the same efficiency constraint, DynaBERT sub-networks consistently outperform other BERT compression methods while matching BERTBASE or RoBERTaBASE with comparable or smaller resource use.

  • Takeaways & Limitations

    Adaptive width and depth provide flexible sub-networks that explore accuracy–size trade-offs across deployment requirements.

  • Takeaways & Limitations

    Easier deployment can intensify negative impacts such as job loss in dialogue systems, and extensions to generative models risk offensive, biased, or unethical outputs.

Abstract

from arXiv · show

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.

1 Introduction

DynaBERT addresses the mismatch between fixed-size BERT compression and edge devices with varying hardware and changing resource availability by enabling flexible architectural configurations. It adapts both width and depth to provide more deployment choices than depth-only approaches.

  • BERT’s parameter, computation, and energy demands hinder deployment on resource-limited edge devices.
  • Different devices and changing workloads require models whose architectures and inference sub-networks can adapt to available resources.
  • Existing compression methods typically produce fixed-size models, while depth-adaptive methods offer limited architectural configurations.
  • DynaBERT jointly adapts BERT width and depth, enabling more architectural configurations and finer accuracy–size trade-offs than depth-only methods.

2 Method

DynaBERT is trained in two distillation stages, first learning adaptive width and then adaptive width plus depth. Its width is controlled through attention heads and FFN neurons, while rewiring prioritizes components shared across sub-networks.

  • Training procedure: Training first distills a fixed teacher into width-adaptive DynaBERTW, then distills DynaBERTW into sub-networks with adaptive width and depth.
  • Training procedure: Skipping the width-adaptive teacher stage or using a depth-adaptive teacher first leads to inferior performance.
  • Adaptive width: Transformer width is adapted by retaining selected attention heads and intermediate FFN neurons, while the embedding dimension remains fixed.
  • Network rewiring: Network rewiring ranks heads and neurons by importance so more important components are shared by more sub-networks.
  • Adaptive-width distillation: Width-stage distillation transfers teacher logits, embeddings, and hidden states to student sub-networks at different widths.
  • Adaptive-depth distillation: Depth-stage distillation uses DynaBERTW at maximum depth as teacher while continuing training across widths and matching hidden states across retained layers.

3 Experiment

Experiments on GLUE and SQuAD evaluate DynaBERT and DynaRoBERTa across adaptive widths, depths, and efficiency constraints. The results show comparable or better performance than base and compressed models, while ablations identify effective training components.

  • Main Results: DynaBERT and DynaRoBERTa achieve comparable performance to BERTBASE and RoBERTaBASE at the same or smaller sizes across GLUE tasks.For most tasks, width and depth can be reduced without performance drop, and the largest sub-network is not always best.
  • Comparison with Other Methods: Under matched efficiency constraints, DynaBERT outperforms DistilBERT and TinyBERT, while DynaRoBERTa outperforms LayerDrop, including LayerDrop trained with more data.The comparison uses parameters, FLOPs, and latency on Nvidia K40 GPU and Kirin 810 ARM CPU.
  • Results on SQuAD: A SQuAD sub-network using only 1/2 the width or depth of BERTBASE achieves comparable or better performance, and DynaBERT outperforms TinyBERT and DistilBERT at matched parameters or FLOPs.The SQuAD evaluation reports EM and F1 on the development set.
  • Ablation Study: Network rewiring raises DynaBERTW average accuracy by over 2 points, while knowledge distillation and data augmentation add a further 1.5 points over the corresponding preceding configuration.The ablation averages accuracy across four width multipliers on GLUE and compares against separate-network and vanilla baselines.
  • Ablation Study: Knowledge distillation and data augmentation significantly improve DynaBERT accuracy over vanilla training on all three evaluated data sets; final fine-tuning helps SST-2 and CoLA but not MRPC.The reported choice uses the higher average validation accuracy before or after fine-tuning.
  • Ablation Study: Using width-adaptive DynaBERTW as a teacher assistant improves DynaBERT on all three investigated data sets, while width-first training outperforms depth-first training at multiplier 0.5.The paper attributes the depth-first degradation to consecutive layer computation and the inability to rewire connections by layer importance.

4 Conclusion

DynaBERT flexibly adjusts model size and latency by selecting sub-networks with different widths and depths, outperforming other BERT compression methods under matched efficiency constraints.

  • DynaBERT flexibly adjusts size and latency by selecting sub-networks with different widths and depths.Width varies attention heads and FFN neurons, while depth varies Transformer layers; knowledge distillation trains the sub-networks.

Broader Impact

DynaBERT targets practical edge deployment by replacing fixed-size compression with adaptive sub-networks, while its broader impacts include both efficiency benefits and risks from easier deployment.

  • DynaBERT adjusts size and latency by selecting sub-networks with adaptive width and depth instead of producing one fixed-size model.The approach supports deployment across hardware platforms with different performance characteristics.
  • One trained DynaBERT can provide different sub-networks for different hardware platforms without separate fine-tuning for each sub-network.A deployed device can also select the same or smaller sub-networks as its efficiency constraints change.
  • Shared weights across sub-networks reduce training and inference costs compared with separately using different-sized models, potentially reducing carbon emissions.The paper describes this as an environmental benefit of the single-model design.
  • The adaptive-width-and-depth training process can act as regularization and may contribute to improved generalization performance.The paper presents this as a possible reason for performance gains at larger model sizes.
  • Easier deployment of BERT can intensify negative impacts, including job loss from dialogue systems and risks of offensive, biased, or unethical outputs from generative extensions.The paper specifically identifies help-desk replacement and extending the method to models such as GPT as examples.

B.1 Description of Data sets in the GLUE benchmark

The GLUE benchmark evaluates diverse natural language understanding abilities, with task-specific metrics and documented training and measurement settings for the experiments.

  • GLUE includes textual entailment, question answering, similarity and paraphrase, sentiment analysis, and linguistic acceptability tasks.The benchmark includes RTE, MNLI, QNLI, MRPC, QQP, STS-B, SST-2, and CoLA; MNLI uses matched and mismatched sections.
  • The GLUE experiments use specified DynaBERT and DynaRoBERTa training hyperparameters, with the same Table 9 settings applied to both model families.The cited appendix passage points to Table 9 for the detailed configurations.
  • SQuAD training separates objectives into two stages because embedding and hidden-state losses are much larger than prediction loss.The first stage uses augmented data for 2 epochs, followed by prediction-loss training on original data for 10 epochs with batch size 12.
  • FLOPs are measured with batch size 1 and sequence length 128, while GPU inference speed uses QNLI with batch size 128 and sequence length 128.Embedding lookup operations are excluded from FLOPs because their inference time is negligible relative to Transformer layers; GPU timing averages 100 batches.

C.1 More Results on the GLUE Benchmark

Across GLUE tasks and efficiency constraints, DynaBERT and DynaRoBERTa match the base models with lower resource use and generally outperform competing compression methods.

  • DynaBERT achieves comparable accuracy to BERTBASE with the same size, while DynaRoBERTa outperforms RoBERTaBASE on seven of eight tasks.The paper attributes the possible gains to the regularization effect of the more difficult adaptive-width-and-depth training.
  • Figure 6 compares DynaBERT and DynaRoBERTa with other compression methods across all GLUE tasks under parameter, FLOPs, GPU-latency, and CPU-latency constraints.The latency settings use an Nvidia K40 GPU and Kirin 810 A76 ARM CPU.
  • On all tasks, DynaBERT and DynaRoBERTa achieve comparable accuracy to BERTBASE and RoBERTaBASE while often requiring fewer parameters, FLOPs, or lower latency.This establishes the main accuracy-efficiency comparison across the evaluated constraints.
  • Under equal parameter constraints, DynaBERT sub-networks outperform DistilBERT except on STS-B and outperform TinyBERT except on MRPC.DynaRoBERTa sub-networks outperform LayerDrop, including LayerDrop trained with substantially more data.

C.2 Full Results of Ablation Study

The ablation results examine how width and depth adaptation, network rewiring, and efficiency constraints affect DynaBERT’s sub-network performance.

  • Training DynaBERTW with Adaptive Width: DynaBERTW performs similarly to the separate-network baseline at its largest width and significantly better at smaller widths.The accuracy gain becomes more significant as width decreases.
  • Training DynaBERTW with Adaptive Width: After network rewiring, average accuracy is over 2 points higher than without rewiring.
  • Training DynaBERT with Adaptive Width and Depth: Table 12 reports accuracy for each combination of width and depth multipliers in the DynaBERT ablation study.
  • Efficiency Comparison: Figure 6 compares parameters, FLOPs, GPU latency, CPU latency, and average MNLI accuracy for DynaBERT, DynaRoBERTa, and other methods.

C.3 Full Results of Different Methods to Train DynaBERTW

The section compares alternative training strategies for DynaBERTW, including progressive rewiring and universally slimmable training, against predefined width-multiplier training.

  • Progressive Rewiring: Progressive rewiring progressively rewires attention heads and neurons as additional width multipliers become supported during training.For multipliers [1.0, 0.75, 0.5, 0.25], rewiring begins with support for [1.0, 0.75] before expanding further.
  • Universally Slimmable Training: Universally slimmable training samples several width multipliers in each training iteration and uses inplace distillation.
  • Reported Results: Table 13 covers progressive rewiring, while Table 14 covers universally slimmable training for DynaBERTW.
  • Comparison Setup: The reported comparison uses width multipliers [1.0, 0.75, 0.5, 0.25] without data augmentation.
  • Comparison Result: Universally slimmable training shows no significant difference from the alternative predefined-multiplier training procedure.

C.4 Looking into DynaBERT

Attention-map analyses inspect DynaBERT’s behavior across widths on CoLA and SST-2, revealing task-dependent patterns and redundancy in later Transformer layers.

  • CoLA: CoLA is a binary single-sentence classification task predicting whether an English sentence is linguistically acceptable.
  • CoLA: For CoLA’s non-acceptable reversed sentence, final-layer attention heads at widths mw = 1.0 and 0.25 attend to words with almost equal probability.
  • CoLA: The CoLA attention behavior is observed in both examined DynaBERT widths and is also found in other CoLA samples.
  • SST-2: On SST-2, most attention maps in the final few layers point to the unused [SEP] token for both mw = 1 and 0.25.
  • SST-2: SST-2 attention patterns indicate redundancy in Transformer layers, consistent with less than 1 point of accuracy degradation at depth multiplier md = 0.5.
  • Attention Maps: Figures 7–9 compare attention maps across DynaBERT widths, BERTBASE, and the CoLA and SST-2 examples described in the text.

E Preliminary Results of Applying DynaBERT in the Pretraining Phase

The preliminary pretraining experiment applies DynaBERT to a six-layer BERT backbone while varying hidden size, attention heads, intermediate neurons, and layer count.

  • Experimental Setup: The experiment uses a pre-trained 6-layer BERT downloaded from the official Google BERT repository as the backbone.
  • Adaptive Dimensions: DynaBERT sub-networks use hidden sizes H = 128, 256, 512, 768 and adjust attention heads, intermediate neurons, and Transformer layers.
  • Pretraining Results: Table 34 reports MNLI-m development accuracy for separately pre-trained BERT models and sub-networks of a pre-trained DynaBERT.
Loading 2004.04037v2…