SuperShaper: Task-Agnostic Super Pre-training of BERT Models with Variable Hidden Dimensions

Abstract

Task-agnostic pre-training followed by task-specific fine-tuning is a default approach to train NLU models. Such models need to be deployed on devices across the cloud and the edge with varying resource and accuracy constraints. For a given task, repeating pre-training and fine-tuning across tens of devices is prohibitively expensive. We propose SuperShaper, a task agnostic pre-training approach which simultaneously pre-trains a large number of Transformer models by varying shapes, i.e., by varying the hidden dimensions across layers. This is enabled by a backbone network with linear bottleneck matrices around each Transformer layer which are sliced to generate differently shaped sub-networks. In spite of its simple design space and efficient implementation, SuperShaper discovers networks that effectively trade-off accuracy and model size: Discovered networks are more accurate than a range of hand-crafted and automatically searched networks on GLUE benchmarks. Further, we find two critical advantages of shape as a design variable for Neural Architecture Search (NAS): (a) heuristics of good shapes can be derived and networks found with these heuristics match and even improve on carefully searched networks across a range of parameter counts, and (b) the latency of networks across multiple CPUs and GPUs are insensitive to the shape and thus enable device-agnostic search. In summary, SuperShaper radically simplifies NAS for language models and discovers networks that generalize across tasks, parameter constraints, and devices.

References in corpus (17)

• Distilling the Knowledge in a Neural Network
• Neural Architecture Search with Reinforcement Learning
• Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer
• DARTS: Differentiable Architecture Search
• The Lottery Ticket Hypothesis: Finding Sparse, Trainable Neural Networks
• Well-Read Students Learn Better: On the Importance of Pre-training Compact Models
• Distilling Task-Specific Knowledge from BERT into Simple Neural Networks
• Reducing Transformer Depth on Demand with Structured Dropout
• A Tensorized Transformer for Language Modeling
• Playing the lottery with rewards and multiple languages: lottery tickets in RL and NLP
• NAS-BERT: Task-Agnostic and Adaptive-Size BERT Compression with Neural Architecture Search
• Are Sixteen Heads Really Better than One?
• BigNAS: Scaling Up Neural Architecture Search with Big Single-Stage Models
• When BERT Plays the Lottery, All Tickets Are Winning
• BERT-of-Theseus: Compressing BERT by Progressive Module Replacing
• EarlyBERT: Efficient BERT Training via Early-bird Lottery Tickets
• You Only Compress Once: Towards Effective and Elastic BERT Compression via Exploit-Explore Stochastic Nature Gradient