Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

AdaBERT: Task-Adaptive BERT Compression with Differentiable Neural Architecture Search

Daoyuan Chen, Yaliang Li, Minghui Qiu, Zhen Wang, Bofang Li, Bolin Ding, Hongbo Deng, Jun Huang, Wei Lin, Jingren Zhou

TL;DR

AdaBERT tackles the challenge of deploying BERT in resource-constrained settings by introducing a task-adaptive compression framework guided by differentiable Neural Architecture Search. It combines a task-oriented knowledge distillation loss with an efficiency-aware constraint to automatically discover compact, task-specific architectures that mimic the teacher while meeting efficiency targets. The approach uses a macro search space, CNN-based operations, Gumbel-Softmax relaxation, and a probe-based knowledge transfer mechanism to tailor models to individual tasks. Empirical results on GLUE show substantial speedups and parameter reductions with only modest drops in accuracy, demonstrating the practical value of task-adaptive BERT compression.

Abstract

Large pre-trained language models such as BERT have shown their effectiveness in various natural language processing tasks. However, the huge parameter size makes them difficult to be deployed in real-time applications that require quick inference with limited resources. Existing methods compress BERT into small models while such compression is task-independent, i.e., the same compressed BERT for all different downstream tasks. Motivated by the necessity and benefits of task-oriented BERT compression, we propose a novel compression method, AdaBERT, that leverages differentiable Neural Architecture Search to automatically compress BERT into task-adaptive small models for specific tasks. We incorporate a task-oriented knowledge distillation loss to provide search hints and an efficiency-aware loss as search constraints, which enables a good trade-off between efficiency and effectiveness for task-adaptive BERT compression. We evaluate AdaBERT on several NLP tasks, and the results demonstrate that those task-adaptive compressed models are 12.7x to 29.3x faster than BERT in inference time and 11.5x to 17.0x smaller in terms of parameter size, while comparable performance is maintained.

AdaBERT: Task-Adaptive BERT Compression with Differentiable Neural Architecture Search

TL;DR

AdaBERT tackles the challenge of deploying BERT in resource-constrained settings by introducing a task-adaptive compression framework guided by differentiable Neural Architecture Search. It combines a task-oriented knowledge distillation loss with an efficiency-aware constraint to automatically discover compact, task-specific architectures that mimic the teacher while meeting efficiency targets. The approach uses a macro search space, CNN-based operations, Gumbel-Softmax relaxation, and a probe-based knowledge transfer mechanism to tailor models to individual tasks. Empirical results on GLUE show substantial speedups and parameter reductions with only modest drops in accuracy, demonstrating the practical value of task-adaptive BERT compression.

Abstract

Large pre-trained language models such as BERT have shown their effectiveness in various natural language processing tasks. However, the huge parameter size makes them difficult to be deployed in real-time applications that require quick inference with limited resources. Existing methods compress BERT into small models while such compression is task-independent, i.e., the same compressed BERT for all different downstream tasks. Motivated by the necessity and benefits of task-oriented BERT compression, we propose a novel compression method, AdaBERT, that leverages differentiable Neural Architecture Search to automatically compress BERT into task-adaptive small models for specific tasks. We incorporate a task-oriented knowledge distillation loss to provide search hints and an efficiency-aware loss as search constraints, which enables a good trade-off between efficiency and effectiveness for task-adaptive BERT compression. We evaluate AdaBERT on several NLP tasks, and the results demonstrate that those task-adaptive compressed models are 12.7x to 29.3x faster than BERT in inference time and 11.5x to 17.0x smaller in terms of parameter size, while comparable performance is maintained.

Paper Structure

This paper contains 26 sections, 6 equations, 3 figures, 5 tables.

Table of Contents

  1. Introduction
  2. Methodology
  3. Overview
  4. Search Space $\mathcal{A}$
  5. Task-oriented Knowledge Distillation
  6. Task-useful Knowledge Probe.
  7. Attentive Hierarchical Transfer.
  8. Efficiency-Aware Loss
  9. Differentiable Architecture Searching
  10. Search Space Setting
  11. Search Algorithm
  12. Experiments
  13. Setup
  14. Datasets.
  15. Baselines.
  16. ...and 11 more sections

Figures (3)

  • Figure 1: The overview of AdaBERT.
  • Figure 2: Search space including stacked layers and stacked cells.
  • Figure 3: The searched basic cells for three kinds of NLP tasks.