# Improving Task-Agnostic BERT Distillation with Layer Mapping Search

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

## Key facts

- Cited by 8 papers in the map (data through 14 Oct 2026).
- First posted on arXiv 11 Dec 2020.

## Built on by (most cited first, 8 of 8)

- [f-Divergence Minimization for Sequence-Level Knowledge Distillation](https://mlsanity.com/paper/2307.15190/) (2023)
- [AutoTinyBERT: Automatic Hyper-parameter Optimization for Efficient Pre-trained Language Models](https://mlsanity.com/paper/2107.13686/) (2021)
- [E-LANG: Energy-Based Joint Inferencing of Super and Swift Language Models](https://mlsanity.com/paper/2203.00748/) (2022)
- [Extract then Distill: Efficient and Effective Task-Agnostic BERT Distillation](https://mlsanity.com/paper/2104.11928/) (2021)
- [SuperShaper: Task-Agnostic Super Pre-training of BERT Models with Variable Hidden Dimensions](https://mlsanity.com/paper/2110.04711/) (2021)
- [Neural Architecture Search for Effective Teacher-Student Knowledge Transfer in Language Models](https://mlsanity.com/paper/2303.09639/) (2023)
- [A Comparative Analysis of Task-Agnostic Distillation Methods for Compressing Transformer Language Models](https://mlsanity.com/paper/2310.08797/) (2023)
- [DistillLens: Symmetric Knowledge Distillation Through Logit Lens](https://mlsanity.com/paper/2602.13567/) (2026)

## Abstract (arXiv)

Knowledge distillation (KD) which transfers the knowledge from a large teacher model to a small student model, has been widely used to compress the BERT model recently. Besides the supervision in the output in the original KD, recent works show that layer-level supervision is crucial to the performance of the student BERT model. However, previous works designed the layer mapping strategy heuristically (e.g., uniform or last-layer), which can lead to inferior performance. In this paper, we propose to use the genetic algorithm (GA) to search for the optimal layer mapping automatically. To accelerate the search process, we further propose a proxy setting where a small portion of the training corpus are sampled for distillation, and three representative tasks are chosen for evaluation. After obtaining the optimal layer mapping, we perform the task-agnostic BERT distillation with it on the whole corpus to build a compact student model, which can be directly fine-tuned on downstream tasks. Comprehensive experiments on the evaluation benchmarks demonstrate that 1) layer mapping strategy has a significant effect on task-agnostic BERT distillation and different layer mappings can result in quite different performances; 2) the optimal layer mapping strategy from the proposed search process consistently outperforms the other heuristic ones; 3) with the optimal layer mapping, our student model achieves state-of-the-art performance on the GLUE tasks.