[Paper Review] AutoBERT-Zero: Evolving BERT Backbone from Scratch

Jiahui Gao, Hang Xu+6 more|arXiv (Cornell University)|Jul 15, 2021

Topic Modeling4 citations

TL;DR

This paper proposes AutoBERT-Zero, a novel neural architecture search framework that evolves a BERT-like backbone from scratch using a flexible search space combining primitive mathematical operations and convolutional blocks. By introducing an Operation-Priority NAS (OP-NAS) algorithm with Bi-branch Weight-Sharing for efficient evaluation, the method discovers a hybrid architecture that outperforms BERT, RoBERTa, and BERT-large on GLUE with 2.4 and 1.4 higher scores respectively, while reducing FLOPs, latency, and memory usage.

ABSTRACT

Transformer-based pre-trained language models like BERT and its variants have recently achieved promising performance in various natural language processing (NLP) tasks. However, the conventional paradigm constructs the backbone by purely stacking the manually designed global self-attention layers, introducing inductive bias and thus leads to sub-optimal. In this work, we make the first attempt to automatically discover novel pre-trained language model (PLM) backbone on a flexible search space containing the most fundamental operations from scratch. Specifically, we propose a well-designed search space which (i) contains primitive math operations in the intra-layer level to explore novel attention structures, and (ii) leverages convolution blocks to be the supplementary for attentions in the inter-layer level to better learn local dependency. To enhance the efficiency for finding promising architectures, we propose an Operation-Priority Neural Architecture Search (OP-NAS) algorithm, which optimizes both the search algorithm and evaluation of candidate models. Specifically, we propose Operation-Priority (OP) evolution strategy to facilitate model search via balancing exploration and exploitation. Furthermore, we design a Bi-branch Weight-Sharing (BIWS) training strategy for fast model evaluation. Extensive experiments show that the searched architecture (named AutoBERT-Zero) significantly outperforms BERT and its variants of different model capacities in various downstream tasks, proving the architecture's transfer and scaling abilities. Remarkably, AutoBERT-Zero-base outperforms RoBERTa-base (using much more data) and BERT-large (with much larger model size) by 2.4 and 1.4 higher score on GLUE test set.

Motivation & Objective

• To overcome the inductive bias and sub-optimal performance of manually designed self-attention mechanisms in BERT and its variants.
• To automatically discover more efficient and powerful attention structures and backbone architectures for pre-trained language models (PLMs) from scratch.
• To design a flexible search space that combines intra-layer operations and inter-layer global (attention) and local (convolution) operations for improved modeling of global and local dependencies.
• To develop an efficient NAS algorithm that reduces computational cost for PLM architecture search, enabling practical discovery of high-performing architectures.
• To validate the transferability and scalability of the searched architecture across different model sizes and downstream NLP tasks.

Proposed method

• Proposes a flexible search space with intra-layer operations (primitive math operations like addition, multiplication, etc.) to explore novel attention mechanisms with variable path lengths and input nodes.
• Introduces inter-layer search space that combines self-attention and convolutional blocks to balance global and local feature learning across layers.
• Develops an Operation-Priority (OP) evolution strategy that uses confidence-based UCB selection to balance exploration and exploitation during architecture mutation.
• Designs a Bi-branch Weight-Sharing (BIWS) training strategy to accelerate model evaluation by sharing weights between super-net branches for attention and convolution components.
• Employs a proxy evaluation pipeline using GLUE finetuning after 40k steps of pre-training to rank candidate architectures efficiently.
• Uses a two-phase training process: first NAS with fast evaluation via BIWS, then full fine-tuning on standard benchmarks with standard hyperparameters.

Experimental results

Research questions

• RQ1 Can a neural architecture search framework discover a more effective and efficient PLM backbone than manually designed self-attention mechanisms?
• RQ2 Does combining learned attention structures with convolutional operations improve performance and efficiency in pre-trained language models?
• RQ3 Can an efficient NAS algorithm significantly reduce the computational cost of searching for PLM backbones while maintaining high performance?
• RQ4 Does the searched architecture generalize well across different model sizes and downstream NLP tasks?
• RQ5 Can the discovered architecture outperform existing state-of-the-art models like RoBERTa and BERT-large with less data and smaller size?

Key findings

• AutoBERT-Zero-base achieves a GLUE score of 86.3, outperforming RoBERTa-base (83.9) and BERT-large (84.9) by 2.4 and 1.4 points respectively.
• The searched architecture reduces FLOPs by 20.7% (2.3e10 vs. 2.9e10) and inference latency by 7.2% (27.2s vs. 29.3s) compared to BERT-base on V100.
• AutoBERT-Zero-base uses 104M parameters, slightly fewer than BERT-base (110M), while achieving better performance.
• The model occupies 2.1G memory, a 16% reduction compared to BERT-base (2.5G), demonstrating improved memory efficiency.
• The BIWS training strategy reduces the total search cost from an estimated 182K GPU hours to 24K GPU hours, significantly improving efficiency.
• The searched architecture demonstrates strong transferability, achieving state-of-the-art performance on GLUE, SQuAD v1.1, SQuAD v2.0, and other benchmarks.