# FlexiBERT: Are Current Transformer Architectures too Homogeneous and Rigid?

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

## Key facts

- Cited by 12 papers in the map (data through 14 Oct 2026).
- Builds on 49 of its 64 references that are in the map.
- First posted on arXiv 23 May 2022.

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

- [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)
- [Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer](https://mlsanity.com/paper/1910.10683/) (2019)
- [RoBERTa: A Robustly Optimized BERT Pretraining Approach](https://mlsanity.com/paper/1907.11692/) (2019)
- [EfficientNet: Rethinking Model Scaling for Convolutional Neural Networks](https://mlsanity.com/paper/1905.11946/) (2019)
- [BART: Denoising Sequence-to-Sequence Pre-training for Natural Language Generation, Translation, and Comprehension](https://mlsanity.com/paper/1910.13461/) (2019)
- [GLUE: A Multi-Task Benchmark and Analysis Platform for Natural Language Understanding](https://mlsanity.com/paper/1804.07461/) (2018)
- [Dropout as a Bayesian Approximation: Representing Model Uncertainty in Deep Learning](https://mlsanity.com/paper/1506.02142/) (2015)
- [How Powerful are Graph Neural Networks?](https://mlsanity.com/paper/1810.00826/) (2018)
- [Distributed Representations of Words and Phrases and their Compositionality](https://mlsanity.com/paper/1310.4546/) (2013)
- [XLNet: Generalized Autoregressive Pretraining for Language Understanding](https://mlsanity.com/paper/1906.08237/) (2019)
- [Optuna: A Next-generation Hyperparameter Optimization Framework](https://mlsanity.com/paper/1907.10902/) (2019)
- [SuperGLUE: A Stickier Benchmark for General-Purpose Language Understanding Systems](https://mlsanity.com/paper/1905.00537/) (2019)
- [Neural Architecture Search with Reinforcement Learning](https://mlsanity.com/paper/1611.01578/) (2016)
- [ShuffleNet: An Extremely Efficient Convolutional Neural Network for Mobile Devices](https://mlsanity.com/paper/1707.01083/) (2017)
- [Self-Attention with Relative Position Representations](https://mlsanity.com/paper/1803.02155/) (2018)
- [Learning Transferable Architectures for Scalable Image Recognition](https://mlsanity.com/paper/1707.07012/) (2017)
- [Cross-lingual Language Model Pretraining](https://mlsanity.com/paper/1901.07291/) (2019)
- [Aligning Books and Movies: Towards Story-like Visual Explanations by Watching Movies and Reading Books](https://mlsanity.com/paper/1506.06724/) (2015)
- [Effective Approaches to Attention-based Neural Machine Translation](https://mlsanity.com/paper/1508.04025/) (2015)
- [MPNet: Masked and Permuted Pre-training for Language Understanding](https://mlsanity.com/paper/2004.09297/) (2020)
- [Regularized Evolution for Image Classifier Architecture Search](https://mlsanity.com/paper/1802.01548/) (2018)
- [Efficient Neural Architecture Search via Parameter Sharing](https://mlsanity.com/paper/1802.03268/) (2018)
- [MobileBERT: a Compact Task-Agnostic BERT for Resource-Limited Devices](https://mlsanity.com/paper/2004.02984/) (2020)
- [AutoML: A Survey of the State-of-the-Art](https://mlsanity.com/paper/1908.00709/) (2019)
- [FNet: Mixing Tokens with Fourier Transforms](https://mlsanity.com/paper/2105.03824/) (2021)
- [NAS-Bench-101: Towards Reproducible Neural Architecture Search](https://mlsanity.com/paper/1902.09635/) (2019)
- [DynaBERT: Dynamic BERT with Adaptive Width and Depth](https://mlsanity.com/paper/2004.04037/) (2020)
- [graph2vec: Learning Distributed Representations of Graphs](https://mlsanity.com/paper/1707.05005/) (2017)
- [Music Transformer](https://mlsanity.com/paper/1809.04281/) (2018)

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

- [AccelTran: A Sparsity-Aware Accelerator for Dynamic Inference with Transformers](https://mlsanity.com/paper/2302.14705/) (2023)
- [Hybrid Quantum Vision Transformers for Event Classification in High Energy Physics](https://mlsanity.com/paper/2402.00776/) (2024)
- [Training-free Neural Architecture Search for RNNs and Transformers](https://mlsanity.com/paper/2306.00288/) (2023)
- [CODEBench: A Neural Architecture and Hardware Accelerator Co-Design Framework](https://mlsanity.com/paper/2212.03965/) (2022)
- [EdgeTran: Co-designing Transformers for Efficient Inference on Mobile Edge Platforms](https://mlsanity.com/paper/2303.13745/) (2023)
- [TransCODE: Co-design of Transformers and Accelerators for Efficient Training and Inference](https://mlsanity.com/paper/2303.14882/) (2023)
- [W-PCA Based Gradient-Free Proxy for Efficient Search of Lightweight Language Models](https://mlsanity.com/paper/2504.15983/) (2025)
- [TabiBERT: A Large-Scale ModernBERT Foundation Model and A Unified Benchmark for Turkish](https://mlsanity.com/paper/2512.23065/) (2025)
- [Compilation and Optimizations for Efficient Machine Learning on Embedded Systems](https://mlsanity.com/paper/2206.03326/) (2022)
- [Zero-Shot Neural Network Evaluation with Sample-Wise Activation Patterns](https://mlsanity.com/paper/2605.07378/) (2026)
- [Designing Compact Neural Architectures via Neuron Gating and Mixed Activation](https://mlsanity.com/paper/2608.14443/) (2026)
- [Neural Spectral Capacity: Measuring and Designing Architectures from Network Specification Alone](https://mlsanity.com/paper/2609.23087/) (2026)

## Abstract (arXiv)

The existence of a plethora of language models makes the problem of selecting the best one for a custom task challenging. Most state-of-the-art methods leverage transformer-based models (e.g., BERT) or their variants. Training such models and exploring their hyperparameter space, however, is computationally expensive. Prior work proposes several neural architecture search (NAS) methods that employ performance predictors (e.g., surrogate models) to address this issue; however, analysis has been limited to homogeneous models that use fixed dimensionality throughout the network. This leads to sub-optimal architectures. To address this limitation, we propose a suite of heterogeneous and flexible models, namely FlexiBERT, that have varied encoder layers with a diverse set of possible operations and different hidden dimensions. For better-posed surrogate modeling in this expanded design space, we propose a new graph-similarity-based embedding scheme. We also propose a novel NAS policy, called BOSHNAS, that leverages this new scheme, Bayesian modeling, and second-order optimization, to quickly train and use a neural surrogate model to converge to the optimal architecture. A comprehensive set of experiments shows that the proposed policy, when applied to the FlexiBERT design space, pushes the performance frontier upwards compared to traditional models. FlexiBERT-Mini, one of our proposed models, has 3% fewer parameters than BERT-Mini and achieves 8.9% higher GLUE score. A FlexiBERT model with equivalent performance as the best homogeneous model achieves 2.6x smaller size. FlexiBERT-Large, another proposed model, achieves state-of-the-art results, outperforming the baseline models by at least 5.7% on the GLUE benchmark.