# NAS-BERT: Task-Agnostic and Adaptive-Size BERT Compression with Neural Architecture Search

arXiv: https://arxiv.org/abs/2105.14444 · HTML: https://mlsanity.com/paper/2105.14444/

## Key facts

- Cited by 56 papers in the map (data through 14 Oct 2026).
- Builds on 45 of its 57 references that are in the map.
- First posted on arXiv 30 May 2021.

## Builds on (most cited first, 30 of 45)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [RoBERTa: A Robustly Optimized BERT Pretraining Approach](https://mlsanity.com/paper/1907.11692/) (2019)
- [GLUE: A Multi-Task Benchmark and Analysis Platform for Natural Language Understanding](https://mlsanity.com/paper/1804.07461/) (2018)
- [DistilBERT, a distilled version of BERT: smaller, faster, cheaper and lighter](https://mlsanity.com/paper/1910.01108/) (2019)
- [SQuAD: 100,000+ Questions for Machine Comprehension of Text](https://mlsanity.com/paper/1606.05250/) (2016)
- [XLNet: Generalized Autoregressive Pretraining for Language Understanding](https://mlsanity.com/paper/1906.08237/) (2019)
- [ALBERT: A Lite BERT for Self-supervised Learning of Language Representations](https://mlsanity.com/paper/1909.11942/) (2019)
- [A Broad-Coverage Challenge Corpus for Sentence Understanding through Inference](https://mlsanity.com/paper/1704.05426/) (2017)
- [Searching for MobileNetV3](https://mlsanity.com/paper/1905.02244/) (2019)
- [DARTS: Differentiable Architecture Search](https://mlsanity.com/paper/1806.09055/) (2018)
- [Know What You Don't Know: Unanswerable Questions for SQuAD](https://mlsanity.com/paper/1806.03822/) (2018)
- [MiniLM: Deep Self-Attention Distillation for Task-Agnostic Compression of Pre-Trained Transformers](https://mlsanity.com/paper/2002.10957/) (2020)
- [TinyBERT: Distilling BERT for Natural Language Understanding](https://mlsanity.com/paper/1909.10351/) (2019)
- [SemEval-2017 Task 1: Semantic Textual Similarity - Multilingual and Cross-lingual Focused Evaluation](https://mlsanity.com/paper/1708.00055/) (2017)
- [Designing Network Design Spaces](https://mlsanity.com/paper/2003.13678/) (2020)
- [Neural Network Acceptability Judgments](https://mlsanity.com/paper/1805.12471/) (2018)
- [Once-for-All: Train One Network and Specialize it for Efficient Deployment](https://mlsanity.com/paper/1908.09791/) (2019)
- [ProxylessNAS: Direct Neural Architecture Search on Target Task and Hardware](https://mlsanity.com/paper/1812.00332/) (2018)
- [Patient Knowledge Distillation for BERT Model Compression](https://mlsanity.com/paper/1908.09355/) (2019)
- [MobileBERT: a Compact Task-Agnostic BERT for Resource-Limited Devices](https://mlsanity.com/paper/2004.02984/) (2020)
- [Reducing Transformer Depth on Demand with Structured Dropout](https://mlsanity.com/paper/1909.11556/) (2019)
- [Single Path One-Shot Neural Architecture Search with Uniform Sampling](https://mlsanity.com/paper/1904.00420/) (2019)
- [Q-BERT: Hessian Based Ultra Low Precision Quantization of BERT](https://mlsanity.com/paper/1909.05840/) (2019)
- [Q8BERT: Quantized 8Bit BERT](https://mlsanity.com/paper/1910.06188/) (2019)
- [DynaBERT: Dynamic BERT with Adaptive Width and Depth](https://mlsanity.com/paper/2004.04037/) (2020)
- [Pay Less Attention with Lightweight and Dynamic Convolutions](https://mlsanity.com/paper/1901.10430/) (2019)
- [BigNAS: Scaling Up Neural Architecture Search with Big Single-Stage Models](https://mlsanity.com/paper/2003.11142/) (2020)
- [Compressing BERT: Studying the Effects of Weight Pruning on Transfer Learning](https://mlsanity.com/paper/2002.08307/) (2020)

## Built on by (most cited first, 30 of 56)

- [SqueezeLLM: Dense-and-Sparse Quantization](https://mlsanity.com/paper/2306.07629/) (2023)
- [MobileLLM: Optimizing Sub-billion Parameter Language Models for On-Device Use Cases](https://mlsanity.com/paper/2402.14905/) (2024)
- [A Survey on Efficient Inference for Large Language Models](https://mlsanity.com/paper/2404.14294/) (2024)
- [Speculative Decoding with Big Little Decoder](https://mlsanity.com/paper/2302.07863/) (2023)
- [A Fast Post-Training Pruning Framework for Transformers](https://mlsanity.com/paper/2204.09656/) (2022)
- [Full Stack Optimization of Transformer Inference: a Survey](https://mlsanity.com/paper/2302.14017/) (2023)
- [Towards Efficient Post-training Quantization of Pre-trained Language Models](https://mlsanity.com/paper/2109.15082/) (2021)
- [A Survey of Techniques for Optimizing Transformer Inference](https://mlsanity.com/paper/2307.07982/) (2023)
- [A Survey of Small Language Models](https://mlsanity.com/paper/2410.20011/) (2024)
- [LLM Performance Predictors are good initializers for Architecture Search](https://mlsanity.com/paper/2310.16712/) (2023)
- [Meta Learning for Natural Language Processing: A Survey](https://mlsanity.com/paper/2205.01500/) (2022)
- [Gradient-Free Structured Pruning with Unlabeled Data](https://mlsanity.com/paper/2303.04185/) (2023)
- [Pruning Self-attentions into Convolutional Layers in Single Path](https://mlsanity.com/paper/2111.11802/) (2021)
- [Analyzing and Mitigating Interference in Neural Architecture Search](https://mlsanity.com/paper/2108.12821/) (2021)
- [EvoPress: Accurate Dynamic Model Compression via Evolutionary Search](https://mlsanity.com/paper/2410.14649/) (2024)
- [EfficientBERT: Progressively Searching Multilayer Perceptron via Warm-up Knowledge Distillation](https://mlsanity.com/paper/2109.07222/) (2021)
- [Exploring Extreme Parameter Compression for Pre-trained Language Models](https://mlsanity.com/paper/2205.10036/) (2022)
- [A Comprehensive Survey of Compression Algorithms for Language Models](https://mlsanity.com/paper/2401.15347/) (2024)
- [Differentially Private Model Compression](https://mlsanity.com/paper/2206.01838/) (2022)
- [FlexiBERT: Are Current Transformer Architectures too Homogeneous and Rigid?](https://mlsanity.com/paper/2205.11656/) (2022)
- [SuperTickets: Drawing Task-Agnostic Lottery Tickets from Supernets via Jointly Architecture Searching and Parameter Pruning](https://mlsanity.com/paper/2207.03677/) (2022)
- [AutoMoE: Heterogeneous Mixture-of-Experts with Adaptive Computation for Efficient Neural Machine Translation](https://mlsanity.com/paper/2210.07535/) (2022)
- [EfficientLLM: Scalable Pruning-Aware Pretraining for Architecture-Agnostic Edge Language Models](https://mlsanity.com/paper/2502.06663/) (2025)
- [Training-free Neural Architecture Search for RNNs and Transformers](https://mlsanity.com/paper/2306.00288/) (2023)
- [AutoDistill: an End-to-End Framework to Explore and Distill Hardware-Efficient Language Models](https://mlsanity.com/paper/2201.08539/) (2022)
- [Mixture-of-Supernets: Improving Weight-Sharing Supernet Training with Architecture-Routed Mixture-of-Experts](https://mlsanity.com/paper/2306.04845/) (2023)
- [Composer: A Search Framework for Hybrid Neural Architecture Design](https://mlsanity.com/paper/2510.00379/) (2025)
- [SuperShaper: Task-Agnostic Super Pre-training of BERT Models with Variable Hidden Dimensions](https://mlsanity.com/paper/2110.04711/) (2021)
- [Constraint-aware and Ranking-distilled Token Pruning for Efficient Transformer Inference](https://mlsanity.com/paper/2306.14393/) (2023)
- [Efficiently Distilling LLMs for Edge Applications](https://mlsanity.com/paper/2404.01353/) (2024)

## Abstract (arXiv)

While pre-trained language models (e.g., BERT) have achieved impressive results on different natural language processing tasks, they have large numbers of parameters and suffer from big computational and memory costs, which make them difficult for real-world deployment. Therefore, model compression is necessary to reduce the computation and memory cost of pre-trained models. In this work, we aim to compress BERT and address the following two challenging practical issues: (1) The compression algorithm should be able to output multiple compressed models with different sizes and latencies, in order to support devices with different memory and latency limitations; (2) The algorithm should be downstream task agnostic, so that the compressed models are generally applicable for different downstream tasks. We leverage techniques in neural architecture search (NAS) and propose NAS-BERT, an efficient method for BERT compression. NAS-BERT trains a big supernet on a search space containing a variety of architectures and outputs multiple compressed models with adaptive sizes and latency. Furthermore, the training of NAS-BERT is conducted on standard self-supervised pre-training tasks (e.g., masked language model) and does not depend on specific downstream tasks. Thus, the compressed models can be used across various downstream tasks. The technical challenge of NAS-BERT is that training a big supernet on the pre-training task is extremely costly. We employ several techniques including block-wise search, search space pruning, and performance approximation to improve search efficiency and accuracy. Extensive experiments on GLUE and SQuAD benchmark datasets demonstrate that NAS-BERT can find lightweight models with better accuracy than previous approaches, and can be directly applied to different downstream tasks with adaptive model sizes for different requirements of memory or latency.