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

Shikhar Tuli<EMAIL_ADDRESS>
Bhishma Dedhia<EMAIL_ADDRESS>
Dept. of Electrical & Computer Engineering, Princeton University
Princeton, NJ 08544 USA Shreshth Tuli<EMAIL_ADDRESS>
Department of Computing, Imperial College London
London, SW7 2AZ UK Niraj K. Jha<EMAIL_ADDRESS>
Dept. of Electrical & Computer Engineering, Princeton University
Princeton, NJ 08544 USA

###### Abstract

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.6$\times$ 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.

##  1 Introduction

In recent years, self-attention (SA)-based transformer models (?, ?) have
achieved state-of-the-art results on tasks that span the natural language
processing (NLP) domain. This burgeoning success has largely been driven by
large-scale pre-training datasets, increasing computational power, and robust
training techniques (?). A challenge that remains is _efficient_ optimal model
selection for a specific task and a set of user requirements. In this context,
only those models should be trained that have the maximum _predicted_
performance. This falls in the domain of neural architecture search (NAS) (?).

###  1.1 Challenges

The design space of transformer models is vast. Several models have been
proposed in the past after rigorous search. Popular models include BERT, XLM,
XLNet, BART, ConvBERT, and FNet (?, ?, ?, ?, ?, ?). Transformer design
involves a choice of several hyperparameters, including the number of layers,
size of hidden embeddings, number of attention heads, and size of the hidden
layer in the feed-forward network (?). This leads to an exponential increase
in the design space, making a brute-force approach to explore the design space
computationally infeasible (?). The aim is to converge to an optimal model as
quickly as possible, by testing the lowest possible number of datapoints (?).
Moreover, model performance may not be deterministic, requiring
heteroscedastic modeling (?).

###  1.2 Existing solutions

Recent NAS advancements use various techniques to explore and optimize
different models in the deep learning domain, from image recognition to speech
recognition and machine translation (?, ?). In the computer-vision domain,
many convolutional neural network (CNN) architectures have been designed using
various search approaches, such as genetic algorithms, reinforcement learning,
structure adaptation, etc. Some even introduce new basic operations (?) to
enhance performance on different tasks. Many works leverage a performance
predictor, often called a surrogate model, to _reliably_ predict model
accuracy. Such a surrogate can be trained through active learning by querying
a few models from the design space and regressing their performance to the
remaining space (under some theoretical assumptions), thus significantly
reducing search times (?, ?).

However, unlike CNN frameworks (?, ?) that are only meant for vision tasks,
there is no universal framework for NLP that differentiates among transformer
architectural hyperparameters. Works that do compare different design
decisions often do not consider _heterogeneity_ and _flexibility_ in their
search space and explore the space over a limited hyperparameter set (?, ?,
?)111Here, by _heterogeneity_ , we mean that different encoder layers can have
distinct attention operations, feed-forward stack depths, etc. By
_flexibility_ , we mean that the hidden dimensions for different encoder
layers, in a transformer architecture, are allowed to be mismatched.. For
instance, Primer (?) only adds depth-wise convolutions to the attention heads;
AutoBERT-Zero (?) lacks deep feed-forward stacks; AutoTinyBERT (?) does not
consider linear transforms (LTs) that have been shown to outperform
traditional SA operations in terms of parameter efficiency; AdaBERT (?) only
considers a design space of convolution and pooling operations. Most works, in
the field of NAS for transformers, target model compression while trying to
maintain the same performance (?, ?, ?), which is orthogonal to our objectives
in this work, _i.e._ , searching for novel architectures that push the
performance frontier. In addition, all previous works only consider _rigid_
architectures. For instance, DynaBERT (?) only adapts the width of the network
by varying the number of attention heads (and _not_ the hidden dimension of
each head), which is only a simple extension to traditional architectures.
Further, their individual models still have the same hidden dimension
throughout the network. AutoTinyBERT (?) and HAT (?), among others, fix the
input and output dimensions for each encoder layer (see Appendix A.1 for a
background on the SA operation), which leads to _rigid_ architectures.

Table 1: Comparison of related works with different parameters (✓indicates
that the corresponding feature is present). Adaptive width refers to different
architectures having possibly different hidden dimensions (albeit each layer
within the architecture having the same hidden dimension). Full flexibility
corresponds to each encoder layer having, possibly, a different hidden
dimension.

Framework | Self-Attention | Conv. | Lin. Transform |  Flexible no. of attn. ops. |  Flexible feed- fwd. stacks | Flexible hidden dim. |  Search technique
---|---|---|---|---|---|---|---
SDP | WMA | DFT | DCT | Ad. width | Full flexibility
Primer |  |  |  |  |  |  |  |  |  |
(?) |  |  | ✓ |  |  | ✓ |  |  |  | ES
AdaBERT |  |  |  |  |  |  |  |  |  |
(?) |  |  | ✓ |  |  |  |  |  |  | DS
AutoTinyBERT |  |  |  |  |  |  |  |  |  |
(?) | ✓ |  |  |  |  | ✓ | ✓ | ✓ |  | ST
DynaBERT |  |  |  |  |  |  |  |  |  |
(?) | ✓ |  |  |  |  | ✓ | ✓ | ✓ |  | ST
NAS-BERT |  |  |  |  |  |  |  |  |  |
(?) | ✓ |  | ✓ |  |  | ✓ |  |  |  | ST
AutoBERT-Zero |  |  |  |  |  |  |  |  |  |
(?) | ✓ |  | ✓ |  |  | ✓ |  |  |  | ES
FlexiBERT (ours) | ✓ | ✓ | ✓ | ✓ | ✓ | ✓ | ✓ | ✓ | ✓ | BOSHNAS

Table 1 gives an overview of various baseline NAS frameworks for transformer
architectures. It presents the aforementioned works and the respective
features they include. Primer (?) and AutoBERT-Zero (?) exploit evolutionary
search (ES), which faces various drawbacks that limit _elitist_ algorithms (?,
?, ?). AdaBERT (?) leverages differentiable architecture search (DS), a
popular technique used in many CNN design spaces (?). On the other hand, some
recent works like AutoTinyBERT (?), DynaBERT (?), and NAS-BERT (?) leverage
super-network training, where one large transformer is trained and its sub-
networks are searched in a _one-shot_ manner. However, this technique is not
amenable to diverse design spaces, as the super-network size would drastically
increase, limiting the gains from weight transfer to the relatively minuscule
sub-network. Moreover, previous works limit their search to either the
standard SA operation, _i.e._ , the scaled dot-product (SDP), or the
convolution operation. We extend the basic attention operation to also include
the weighted multiplicative attention (WMA). Taking motivation from recent
advances with LT-based transformer models (?), we also add discrete Fourier
transform (DFT) and discrete cosine transform (DCT) to our design space.
AutoTinyBERT and DynaBERT also allow adaptive widths in the transformer
architectures in their design space, however, each instance still has the same
dimensionality throughout the network (in other words, every encoder layer has
the same hidden dimension, as explained above). We detail why this is
inherently a limitation in traditional transformer architectures in Appendix
A.1. FlexiBERT, to the best of our knowledge, is the first framework to allow
full _flexibility_ – not only different transformer instances in the design
space can have distinct widths, but each encoder layer within a transformer
instance can also have different hidden dimensions. Finally, we also leverage
a novel NAS technique – Bayesian Optimization using Second-Order Gradients and
Heteroscedastic Models for Neural Architecture Search (BOSHNAS).

###  1.3 Our contributions

To address the limitations of _homogeneous_ and _rigid_ models, we make the
following technical contributions:

  * •

We expand the design space of transformer hyperparameters to incorporate
heterogeneous architectures that venture beyond simple SA by employing other
operations like convolutions and LTs.

  * •

We propose novel projection layers and _relative_ /_trained_ positional
encodings to make hidden sizes _flexible_ across layers – hence the name
FlexiBERT.

  * •

We propose Transformer2vec that uses similarity measures to compare
computational graphs of transformer models to obtain a dense embedding that
captures model similarity in a Euclidean space.

  * •

We propose a novel NAS framework, namely, BOSHNAS. It uses a neural network as
a heteroscedastic surrogate model and second-order gradient-based optimization
using backpropagation to input (GOBI) (?) to speed up search for the next
query in the exploration process. It leverages _nearby_ trained models to
transfer weights in order to reduce the amortized search time for every query.

  * •

Experiments on the GLUE benchmark (?) show that BOSHNAS applied to the
FlexiBERT design space results in a score improvement of 0.4% compared to the
baseline, _i.e._ , NAS-BERT (?). The proposed model, FlexiBERT-Mini, has 3%
fewer parameters than BERT-Mini and achieves 8.9% higher GLUE score. FlexiBERT
also outperforms the best _homogeneous_ architecture by 3%, while requiring
2.6$\times$ fewer parameters. FlexiBERT-Large, our BERT-Large (?) counterpart,
outperforms the state-of-the-art models by at least 5.7% average accuracy on
the first eight tasks in the GLUE benchmark (?).

The rest of the paper is organized as follows. Section 2 presents related
work. Section 3 describes the set of steps and decisions that undergird the
FlexiBERT framework. In Section 4, we present the results of design space
exploration experiments. Finally, Section 5 concludes the article.

##  2 Related Work

We briefly describe related work next.

###  2.1 Transformer design space

Traditionally, transformers have primarily relied on the SA operation (details
in Appendix A.1). Nevertheless, several works have proposed various compute
blocks to reduce the number of model parameters and hence computational cost,
without compromising performance. For instance, ConvBERT uses dynamic span-
based convolutional operations that replace SA heads to directly model local
dependencies (?). Recently, FNet improved model efficiency using LTs instead
(?). MobileBERT, another recent architecture, uses bottleneck structures and
multiple feed-forward stacks to obtain smaller and faster models while
achieving competitive results on well-known benchmarks (?). For completeness,
we present other previously proposed advances to improve the BERT model in
Appendix A.2.

###  2.2 Neural architecture search

NAS is an important machine learning technique that algorithmically searches
for new neural network architectures within a pre-specified design space under
a given objective (?). Prior work has implemented NAS using a variety of
techniques, albeit limited to the CNN design space. A popular approach is to
use a reinforcement learning algorithm, REINFORCE, that has been shown to be
superior to other tabular approaches (?). Other approaches include Gaussian-
Process-based Bayesian Optimization (GP-BO) (?), Evolutionary Search (ES) (?,
?), etc. However, these methods come with challenges that limit their ability
to reach state-of-the-art results in the CNN design space (?).

Recently, NAS has also seen application of surrogate models for performance
prediction in CNNs (?). This results in training of much fewer models to
predict accuracy for the entire design space under some confidence
constraints. However, these predictors are computationally expensive to train.
This leads to a bottleneck, especially in large design spaces, in the training
of subsequent models since new queries are produced only after this predictor
is trained for every batch of trained models in the search space. ? (?) use a
Graph Isomorphism Net (?) that regresses performance values directly on the
computational graphs formed for each CNN model.

Although previously restricted to CNNs (?), NAS has recently seen applications
in the transformer space as well. ? (?) use standard NAS techniques to search
for optimal transformer architectures. However, their method requires every
new model to be trained from scratch. They do not employ knowledge transfer,
in which weights from previously trained _neighboring_ models are used to
speed up subsequent training. This is important in the transformer space since
pre-training every model is computationally expensive. Further, the attention
heads in their model follow the same dimensionality, _i.e._ , are not fully
_flexible_.

One of the state-of-the-art NAS techniques, BANANAS, implements Bayesian
Optimization (BO) over a neural network model and predicts performance
uncertainty using ensemble networks that are, however, too compute-heavy (?).
BANANAS uses mutation/crossover on the current set of best-performing models
and obtains the next best predicted model in this local space. Instead, we
propose the use of GOBI (?) in order to efficiently search for the next query
in the _global_ space. Thanks to random cold restarts, GOBI can search over
diverse models in the architecture space. BANANAS also uses path embeddings,
which have been shown to perform sub-optimally for search over a diverse space
(?).

###  2.3 Graph Embeddings that Drive NAS

Many works on NAS for CNNs have primarily used graph embeddings to model their
performance predictor. These embeddings are formed for each _computational
graph_ , representing a specific CNN architecture in the design space. A
popular approach to learning with graph-structured data is to make use of
graph kernel functions that measure similarity between graphs. A recent work,
NASGEM (?), uses the Weisfeiler-Lehman (WL) sub-tree kernel, which compares
tree-like substructures of two computational graphs. This helps distinguish
between substructures that other kernels, like the random walk kernel, may
deem identical (?). Also, the WL kernel has an attractive computational
complexity. This has made it one of the most widely used graph kernels. Graph-
distance-driven NAS often leads to enhanced representation capacity that
yields optimal search results (?). However, the WL kernel only computes sub-
graph similarities based on overlap in graph nodes. It does not consider
whether two nodes are _inherently_ similar or not. For example, a
computational ‘block’ (or its respective graph node) for an SA head with
$h=128$ and $o=$ SDP, would be closer to another attention block with, say,
$h=256$ and $o=$ WMA, but would be farther from a block representing a feed-
forward layer (for details on SA types, see Section 3.1).

Once we have similarities computed between every possible graph pair in the
design space, next we learn dense embeddings, the Euclidean distance for which
should follow the similarity function. These embeddings would not only be
helpful in effective visualization of the design space, but also for fast
computation of _neighboring_ graphs in the active-learning loop. Further, a
dense embedding helps us practically train a finite-input surrogate function
(as opposed to the sparse path encodings used in ?). Many works have achieved
this using different techniques. ? (?) train task-specific graph embeddings
using a skip-gram model and negative sampling, taking inspiration from
word2vec (?). In this work, we take inspiration from GloVe instead (?), by
applying manifold learning to all distance pairs (?). Hence, using global
similarity distances built over domain knowledge, and batched gradient-based
training, we obtain the proposed Transformer2vec embeddings that are superior
to traditional generalized graph embeddings.

We take motivation from NASGEM (?), which showed that training a WL kernel-
guided encoder has advantages in scalable and flexible search. Thus, we train
a performance predictor on the Transformer2vec embeddings, which not only aid
in transfer of weights between neighboring models, but also support better-
posed continuous performance approximation. More details on the computation of
these embedding are given in Section 3.3.

##  3 Methodology

Figure 1: Overview of the FlexiBERT pipeline.

In this work, we train a heteroscedastic surrogate model that predicts the
performance of a transformer architectrue and uses it to run second-order
optimization in the design space. We do this by decoupling the training
procedure from pre-processing of the embedding of every model in the design
space to speed up training. First, we train embeddings to map the space of
computational graphs to a Euclidean space (Transformer2vec) and then train the
surrogate model on the embeddings.

Our work involves exploring a vast and heterogeneous design space and
searching for optimal architectures with a given task. To this end, we (a)
define a design space via a flexible set of architectural choices (see Section
3.1), (b) generate possible computational graphs ($\mathcal{G}$; see Section
3.2), (c) learn an embedding for each point in the space using a distance
metric for graphs ($\Delta$; see Section 3.3), and (d) employ a novel search
technique (BOSHNAS) based on surrogate modeling of the performance and its
uncertainty over the continuous embedding space (see Section 3.4). In
addition, to tackle the enormous design space, we propose a hierarchical
search technique that iteratively searches over finer-grained models derived
from (e) a crossover of the best models obtained in the current iteration and
their neighbors. Figure 1 gives a broad overview of the FlexiBERT pipeline as
explained above. An unrolled version of this iterative flow is presented
below:

$\text{Design
Space}\rightarrow\mathcal{G}_{1}\xrightarrow{\mathcal{T}}\Delta_{1}\xrightarrow{\text{BOSHNAS}}g^{*}\xrightarrow{\text{cross-
over}}\mathcal{G}_{2}\xrightarrow{\mathcal{T}}\ldots$

However, for simplicity of notation, we omit the iteration index in further
references. We now discuss the key elements of this pipeline in detail.

###  3.1 FlexiBERT Design Space

We now describe the FlexiBERT design space, i.e., box (a) in Figure 1.

####  3.1.1 Set of operations in FlexiBERT

The traditional BERT model is composed of multiple layers, each containing a
bidirectional multi-headed by SA module followed by a feed-forward module.
Several modifications have been proposed to the original encoder, primarily to
the attention module. This gives rise to a richer design space. We consider
WMA-based SA in addition to SDP-based operations (?).

We also incorporate LT-based attention in FNet (?) and dynamic-span-based
convolution (DSC) in ConvBERT (?), in place of the vanilla SA mechanism.
Whereas the original FNet implementation uses DFT, we also consider DCT. The
motivation behind using DCT is its widespread application in lossy data
compression, which we believe can lead to sparse weights, thus leaving room
for optimizations with sparsity-aware machine learning accelerators (?). Our
design space allows variable kernel sizes for convolution-based attention.
Consolidating different attention module types that vary in their
computational costs into a single design space enables the models to have
inter-layer variance in expression capacity. Inspired by MobileBERT (?), we
also consider architectures with multiple feed-forward stacks. We summarize
the entire design space with the range of each operation type in Table 2. The
ranges of different hyperparameters are in accordance with the design space
spanned by BERT-Tiny to BERT-Mini (?), with additional modules included as
discussed. We call this the Tiny-to-Mini space. This restricts our curated
testbed to models with up to $3.3\times 10^{7}$ trainable parameters. This
curated parameter space allows us to perform extensive experiments, comparing
the proposed approach against various baselines (see Section 4.3).

Table 2: Design space description. Super-script ($j$) depicts the value for
layer $j$.

Design Element | Allowed Values
---|---
Number of encoder layers ($l$) | $\\{2,4\\}$
Type of attention operation used ($o^{j}$) | {SA, LT, DSC}
Number of operation heads ($n^{j}$) | $\\{2,4\\}$
Hidden size ($h^{j}$) | $\\{128,256\\}$
Feed-forward dimension ($f^{j}$) | $\\{512,1024\\}$
Number of feed-forward stacks | $\\{1,3\\}$
Operation parameters ($p^{j}$): |
if $o^{j}=$ SA | Self-attention type: {SDP, WMA}
else if $o^{j}=$ LT | Linear transform type: {DFT , DCT}
else if $o^{j}=$ DSC | Convolution kernel size: $\\{5,9\\}$

Every model in the design space can therefore be expressed via a model card
that is a dictionary containing the chosen value for each design decision.
BERT-Tiny (?), in this formulation, can be represented as

$\displaystyle\Big{\\{}l:2,o:[\text{SA},\text{SA}],h:[128,128],n:[2,2],f:\left[[512],[512]\right],p:[\text{SDP},\text{SDP}]\Big{\\}}.$

where the length of the list for every entry in $f$ denotes the size of the
feed-forward stack. The model card can be used to derive the computational
graph of the model using smaller modules inferred from the design choice
(details in Section 3.2).

####  3.1.2 Flexible hidden dimensions

In traditional transformer architectures, the flow of information is
restricted through the use of a constant embedding dimension across the
network (a matrix of dimensions $N_{T}\times h$ from one layer to the next,
where $N_{T}$ denotes the number of tokens and $h$ the hidden dimension; more
details in Appendix A.1). We allow architectures in our design space to have
_flexible_ dimensions across layers that, in turn, enables different layers to
capture information of different dimensions, as it learns more abstract
features deeper into the network. For this, we make the following
modifications:

  * •

Projection layers: We add an affine projection network between encoder layers
with dissimilar hidden sizes to transform encoding dimensionality.

  * •

Relative positional encoding: The vanilla-BERT implementation uses an absolute
positional encoding at the input and propagates it ahead through residual
connections. Since we relax the restriction of a constant hidden size across
layers, this is not applicable to many models in our design space (as the
learned projections for absolute encodings may not be one-to-one). Instead, we
add a _relative_ positional encoding at each layer (?, ?, ?). Such an encoding
can entirely replace absolute positional encodings with relative position
representations learned using the SA mechanism. Whereas the SA module
implementation remains the same as in previous works, for DSC-based and LT-
based attention, we learn the relative encodings separately using SA and add
them to the output of the attention module.

Formally, let $\mathbf{Q}$ and $\mathbf{V}$ denote the query and the value
layers, respectively. Let $\mathbf{R}$ denote the relative embedding tensor
that is to be learned. Let $\mathbf{Z}$ and $\mathbf{X}$ denote the output and
the input tensors of the attention module, respectively. In addition, let us
define LT-based attention and DSC-based attention as $\text{LT}(\cdot)$ and
$\text{DSC}(\cdot)$, respectively. Then,

$\mathrm{RelativeAttention}(\mathbf{X})=\mathrm{softmax}\left(\frac{\mathbf{QR}^{\top}}{\sqrt{d_{\mathbf{Q}}}}\right)\mathbf{V}$
$\displaystyle\mathbf{Z}_{\text{LT}}$
$\displaystyle=\text{LT}(\mathbf{X})+\mathrm{RelativeAttention}(\mathbf{X})$
$\displaystyle\mathbf{Z}_{\text{DSC}}$
$\displaystyle=\text{DSC}(\mathbf{X})+\mathrm{RelativeAttention}(\mathbf{X})$

It should be noted that the proposed approach would only be applicable when
the positional encodings are _trained_ , instead of being predetermined (?).
Thanks to _relative_ and _trained_ positional encodings, this enabled us to
make the dimensionality of data flow _flexible_ across the network layers.
This also means that each layer in the feed-forward stack can have a distinct
hidden dimension.

###  3.2 Graph Library

We now describe the graph library, i.e., box (b) in Figure 1.

####  3.2.1 Block-level computational graphs

To learn a lower-dimensional dense manifold of the given design space,
characterized by a large number of FlexiBERT models, we convert each model
into a computational graph. This is formulated based on the forward flow of
connections for each compute block. For our design space, we take all possible
combinations of the compute blocks derived from the design decisions presented
in Table 2 (see Appendix B.1).

Using this design space and the set of compute blocks, we create all possible
computational graphs within the design space for every transformer model. We
then use recursive hashing as follows (?). For every node in this graph, we
concatenate the hash of its input, hash of that node, and its output, and then
take the hash of the result. We use SHA256 as our hashing function. Doing this
for all nodes and then hashing the concatenated hashes gives us the resultant
hash of a given computational graph. This helps us detect isomorphic graphs
and remove redundancy.

Figure 2: Block-level computation graph for BERT-Tiny in FlexiBERT. The
projection layer implements an identity function since the hidden sizes of
input and output encoder layers are equal.

Figure 2 shows the block-level computational graph for BERT-Tiny. Using the
connection patterns for every possible block permutation (as presented in
Appendix B.1), we can generate multiple graphs for the given design space.

####  3.2.2 Levels of hierarchy

The total number of possible graphs in the design space with _heterogeneous_
feed-forward hidden layers is $\sim$3.32 billion. This is substantially larger
than any transformer design space used in the past.

To make our approach tractable, we propose a hierarchical search method. Each
model in the design space can be considered to be composed of multiple stacks,
where each stack contains at least one encoder layer. In the first step, we
restrict each stack to $s=2$ layers, where each layer in a stack shares the
same design configuration. Naturally, this limits the search space size (the
set of all graphs in this space is denoted by $\mathcal{G}_{1}$). Hence, for
instance, BERT-Tiny falls under $\mathcal{G}_{1}$ since the two encoder layers
have the same configuration. We learn embeddings in this space and then run
NAS to obtain the best-performing models. In the subsequent step, we consider
a design space constituted by a finer-grained neighborhood of these models.
The neighborhood is derived by pairwise crossover between the best-performing
models and their neighbors in a space where the number of layers per stack is
set to $s/2=1$, denoted by $\mathcal{G}_{2}$ (details in Appendix B.3).
Finally, we include _heterogeneous_ feed-forward stacks ($s=1^{*}$) and denote
the space by $\mathcal{G}_{3}$.

###  3.3 Transformer2vec

We now describe the Transformer2vec embedding and how we create an embedding
library from a graph library $\mathcal{G}$, i.e., box (c) in Figure 1.

####  3.3.1 Graph edit distance

Taking inspiration from ? (?) and ? (?), we train dense embeddings using
_global_ distance metrics, such as the Graph Edit Distance (GED) (?). These
embeddings enable fast derivation of _neighboring_ graphs in the active
learning loop to facilitate transfer of weights. We call them Transformer2vec
embeddings. Unlike other approaches like the WL kernel, GED bakes in _domain
knowledge_ in graph comparisons, as explained in Section 2.3, by using a
weighted sum of node insertions, deletions, and substitutions costs.

For the GED computation, we first sort all possible compute blocks in the
order of their computational complexity. Then, we weight the insertion and
deletion cost for every block based on its index in this sorted list, and the
substitution cost between two blocks based on the difference in the indices in
this sorted list. For computing the GED, we use a depth-first algorithm that
requires less memory than traditional methods (?).

####  3.3.2 Training embeddings

Given that there are $S$ graphs in $\mathcal{G}$, we compute the GED for all
possible computational graph pairs. This gives us a dataset of $N={S\choose
2}$ distances. To train the embedding, we minimize the mean-square error as
the loss function between the predicted Euclidean distance and the
corresponding GED. For the design space in consideration, embeddings of $d$
dimensions are generated for every level of the hierarchy. Concretely, to
train embedding $\mathcal{T}$, we minimize the loss

$\displaystyle\mathcal{L}_{\mathcal{T}}=\sum_{1\leq i\leq N,1\leq j\leq
N,i\neq j}\
\Big{(}\mathbf{d}(\mathcal{T}({g_{i}}),\mathcal{T}({g_{j}}))-\mathrm{GED}(g_{i},g_{j})\Big{)}^{2},$

where $\mathbf{d}(\cdot,\cdot)$ is the Euclidean distance and the
$\mathrm{GED}$ is calculated for the corresponding computational graphs
$g_{i}$, $g_{j}\in\mathcal{G}$.

####  3.3.3 Weight transfer among neighboring models

Pre-training each model in the design space is computationally expensive.
Hence, we rely on weight sharing to initialize a query model in order to
_directly_ fine-tune it and minimize exploration time (details in Appendix
B.2). We do this by generating $k$ nearest neighbors of a graph in the design
space (we use $k=100$ for our experiments). Naturally, we would like to
transfer weights from the corresponding fine-tuned neighbor that is closest to
the query, as such models intuitively have similar initial internal
representations.

We calculate this similarity using a biased overlap measure that counts the
number of encoder layers from the input to the output that are common to the
current graph (_i.e._ , have exactly the same set of hyperparameter values).
We stop counting the overlap when we encounter different encoder layers,
regardless of subsequent overlaps. In this ranking, there could be more than
one graph with the same biased overlap with the current graph. Since the
internal representations learned depend on the subsequent set of operations as
well, we break ties based on the embedding distance of these graphs with the
current graph. This gives us a set of neighbors, denoted by $N_{q}$ for a
model $q$, for every graph that are ranked based on both the biased overlap
and the embedding distance. It helps increase the probability of finding a
trained neighbor with high overlap.

As a hard constraint, we only consider transferring weights if the _biased
overlap fraction_ ($\mathcal{O}_{f}(q,n)=\text{\emph{biased overlap}}/l_{q}$,
where $q$ is the query model, $n\in N_{q}$ is the neighbor in consideration,
and $l_{q}$ is the number of layers in $q$) between the queried model and its
neighbor is above a threshold $\tau$. If the constraint is met, the weights of
the shared part from the corresponding neighbor are transferred to the query
and the query is fine-tuned. Otherwise, we pre-train the query. The weight
transfer operation is denoted by $W_{q}\leftarrow W_{n}$.

###  3.4 BOSHNAS

We now describe the BOSHNAS search policy, i.e., box (d) in Figure 1.

####  3.4.1 Uncertainty types

To overcome the challenges of an unexplored design space, it is important to
consider the uncertainty in model predictions to guide the search process.
Predicting model performance deterministically is not enough to estimate the
next most probably best-performing model. We leverage the upper confidence
bound (UCB) exploration on the predicted performance of unexplored models (?).
This could arise from not only the approximations in the surrogate modeling
process, but also parameter initializations and variations in model
performance due to different training recipes. These are called _epistemic_
and _aleatoric_ uncertainties, respectively. The former, also called reducible
uncertainty, arises from a lack of knowledge or information, and the latter,
also called irreducible uncertainty, refers to the inherent variation in the
system to be modeled.

####  3.4.2 Surrogate model

In BOSHNAS, we use Monte-Carlo (MC) dropout (?) and a Natural Parameter
Network (NPN) (?) to model the epistemic and aleatoric uncertainties,
respectively. The NPN not only helps with a distinct prediction of aleatoric
uncertainty that can be used for optimizing the training recipe once we are
close to the optimal architecture, but also serves as a superior model than
Gaussian Processes, Bayesian Neural Networks (BNNs), and other Fully-Connected
Neural Networks (FCNNs) (?). Consider the NPN network $f_{S}(x;\theta)$ with a
transformer embedding $x$ as an input and parameters $\theta$. The output of
such a network is the pair $(\mu,\sigma)\leftarrow f_{S}(x;\theta)$, where
$\mu$ is the predicted mean performance and $\sigma$ is the aleatoric
uncertainty. To model the epistemic uncertainty, we use two deep surrogate
models: (1) teacher ($g_{S}$) and (2) student ($h_{S}$) networks. It is a
surrogate for the performance of a transformer, using its embedding $x$ as an
input. The teacher network is an FCNN with MC Dropout (parameters
$\theta^{\prime}$). To compute the epistemic uncertainty, we generate $n$
samples using $g_{S}(x,\theta^{\prime})$. The standard deviation of the sample
set is denoted by $\xi$. To run GOBI (?) and avoid numerical gradients due to
their poor performance, we use a student network (FCNN with parameters
$\theta^{\prime\prime}$) that directly predicts the output
$\hat{\xi}\leftarrow h_{S}(x,\theta^{\prime\prime})$, a surrogate of $\xi$.

####  3.4.3 Active learning and optimization

For a design space $\mathcal{G}$, we first form an embedding space $\Delta$ by
transforming all graphs in $\mathcal{G}$ using the Transformer2vec embedding.
Assuming we have the three networks $f_{S},g_{S}$, and $h_{S}$, for our
surrogate model, we use the following UCB estimate:

$\displaystyle\begin{split}\mathrm{UCB}&=\mu+k_{1}\cdot\sigma+k_{2}\cdot\hat{\xi}=\Big{(}f_{S}(x,\theta)[0]+k_{1}\cdot
f_{S}(x;\theta)[1]\Big{)}+k_{2}\cdot
h_{S}(x,\theta^{\prime\prime}),\end{split}$ (1)

where $x\in\Delta$, $k_{1}$, and $k_{2}$ are hyperparameters.

To generate the next transformer to test, we run GOBI using neural network
inversion and the AdaHessian optimizer (?) that uses second-order updates to
$x$ ($\nabla^{2}_{x}\mathrm{UCB}$) up till convergence. From this, we get a
new query embedding, $x^{\prime}$. The nearest transformer architecture is
found based on the Euclidean distance of all available transformer
architectures in the design space $\Delta$, giving the next closest model $x$.
This model is fine-tuned (or pre-trained if there is no nearby trained model
with sufficient overlap; see Section 3.3) on the required task to give the
respective performance. Once the new datapoint, $(x,o)$, is received, we train
the models using the loss functions on the updated corpus, $\delta^{\prime}$:

$\displaystyle\mathcal{L}\textsubscript{NPN}(f_{S},x,o)$
$\displaystyle=\sum_{(x,o)\in\delta^{\prime}}\frac{(\mu-o)^{2}}{2\sigma^{2}}+\frac{1}{2}\ln{\sigma^{2}},$
(2) $\displaystyle\mathcal{L}\textsubscript{Teacher}(g_{S},x,o)$
$\displaystyle=\sum_{(x,o)\in\delta^{\prime}}(g_{S}(x,\theta^{\prime})-o)^{2},$
$\displaystyle\mathcal{L}\textsubscript{Student}(h_{S},x)$
$\displaystyle=\sum_{x,\forall(x,o)\in\delta^{\prime}}(h_{S}(x,\theta^{\prime\prime})-\xi)^{2},$

where $\mu,\sigma$ = $f_{S}(x,\theta)$, and $\xi$ is obtained by sampling
$g_{S}(x,\theta^{\prime})$. The first is the aleatoric loss to train the NPN
model (?); the other two are squared-error loss functions. Appendix B.4
presents the flow of these models in a schematic. We run multiple random cold
restarts of GOBI to get multiple queries for the next step in the search
process.

Result: best architecture

1 Initialize: overlap threshold ($\tau$), convergence criterion, uncertainty
sampling prob. ($\alpha$), diversity sampling prob. ($\beta$), surrogate model
($f_{S}$, $g_{S}$, and $h_{S}$) on initial corpus $\delta$, design space
$g\in\mathcal{G}\Leftrightarrow x\in\Delta$;

2while  _convergence criterion not met_ do

3 wait till a worker is free

4 if  _prob $\sim U(0,1)<1-\alpha-\beta$_ then

5 $\delta\leftarrow\delta\ \cup$ {new performance point ($x,o$)};

6 fit(surrogate, $\delta$) using Eqn. (2);

$x$ $\leftarrow$ GOBI($f_{S}$, $h_{S}$);

/* Optimization step */

7 for  _$n$ in $N_{x}$_ do

8 if  _$n$ is trained & $\mathcal{O}_{f}(x,n)\geq\tau$_ then

9 $W_{x}\leftarrow W_{n}$;

10 send $x$ to worker;

11 break;

12 else

13 if  _$1-\alpha-\beta\leq$ prob. $<1-\beta$_ then

$x$ $\leftarrow$ argmax($k_{1}\cdot\sigma+k_{2}\cdot\hat{\xi}$);

/* Uncertainty sampling */

14 send $x$ to worker;

15 else

send random $x$ to worker;

/* Diversity sampling */

16

Algorithm 1 BOSHNAS

Algorithm 1 summarizes the BOSHNAS workflow. Starting from an initial pre-
trained set $\delta$ in the first level of hierarchy $\mathcal{G}_{1}$, we run
until convergence the following steps in a multi-worker compute cluster. To
trade off between exploration and exploitation, we consider two probabilities:
uncertainty-based exploration ($\alpha$) and diversity-based exploration
($\beta$). With probability $1-\alpha-\beta$, we run second-order GOBI using
the surrogate model to minimize UCB in Eq. (1). Adding the converged point
$(x,o)$ in $\delta$, we minimize the loss values in Eq. (2) (line 1 in
Algorithm 1). We then generate a new query point, transfer weights from a
neighboring model, and train it (lines 1-1). With $\alpha$ probability, we
sample the search space using the combination of aleatoric and epistemic
uncertainties, $k_{1}\cdot\sigma+k_{2}\cdot\hat{\xi}$, to find a point where
the performance estimate is uncertain (line 1). To avoid getting stuck in a
localized search subset, we also choose a random point with probability
$\beta$ (line 1). Once we converge in the first level, we continue with second
and third levels, $\mathcal{G}_{2}$ and $\mathcal{G}_{3}$, as described in
Section 3.2.

##  4 Experimental Results

In this section, we show how the FlexiBERT model obtained from BOSHNAS
outperforms the baselines.

###  4.1 Setup

For our experiments, we set the number of layers in each stack to $s=2$ for
the first level of the hierarchy, where models have the same configurations in
every stack. In the second level, we use $s=1$. Finally, we also make the
feed-forward stacks heterogeneous ($s=1^{*}$) in the third level (details
given in Section 3.2). For the range of design choices in Table 2 and setting
$s=2$, we obtained 9312 unique graphs after removing isomorphic graphs. The
dimension of the Transformer2vec embedding is set to $d=16$ after running grid
search. To do this, we minimize the distance prediction error while also
keeping $d$ small using knee-point detection. The hyperparameter values in
Algorithm 1 are obtained through grid search. We use overlap threshold
$\tau=80\%$, $\alpha=\beta=0.1$, and $k_{1}=k_{2}=0.5$ in our experiments. The
convergence criterion is met in BOSHNAS when the change in performance is
within $10^{-4}$ for five iterations. Further experimental setup details are
given in Appendix B.5.

Figure 3: Bar plot comparing all NAS techniques with (a) naive embeddings and
a design space of homogeneous models, (b) naive embeddings and an expanded
design space of homogeneous and heterogeneous models, and (c) Transformer2vec
(T2v) embeddings with the expanded design space. Plotted with 90% confidence
intervals.

###  4.2 Pre-training and Fine-tuning Models

Our pre-training recipe is adapted from the one used in RoBERTa, proposed by ?
(?), with slight variations in order to reduce the training budget (details in
Appendix B.5).

We initialize the architecture space with models adapted from equivalent
models presented in literature (?, ?, ?). The 12 initial models used to
initiate the search process are BERT-Tiny, BERT-2/256 (with two encoder layers
and a fixed hidden dimension of 256), BERT-4/128, BERT-Mini, FNet-Tiny,
FNet-2/256, FNet-4/128, FNet-Mini, ConvBERT-Tiny, ConvBERT-2/256,
ConvBERT-4/128, and ConvBERT-Mini (with $p^{j}=$ DFT for FNets and $p^{j}=9$
for ConvBERTs adapted from the original models). These models form the initial
set $\delta$ in Algorithm 1.

Figure 4: Performance results: (a) best GLUE score with trained models for NAS
baselines and (b) ablation of BOSHNAS. Plotted with 90% confidence intervals.

###  4.3 Ablation Study of BOSHNAS

We compare BOSHNAS against other popular techniques from the CNN space, namely
Random Search (RS), ES, REINFORCE, GP-BO, and a recent state-of-the-art,
BANANAS. We present performance on the GLUE benchmark.

Figure 3 presents the best GLUE scores reached by respective baseline NAS
techniques along with BOSHNAS used with naive (_i.e._ , feature-based one-hot)
or Transformer2vec embeddings on a representative design space. We use the
space in the first level of the hierarchy (_i.e._ , with 9312 graphs, $s=2$)
and run all these algorithms in an active-learning scenario (all targeted
homogeneous models form a subset of this space) over 50 runs for each
algorithm. The plot highlights that enhancing the richness of the design space
enables the algorithms to search for more accurate models (6% improvement
averaged across all models). We also see that Transformer2vec embeddings help
NAS algorithms reach better-performing architectures (9% average improvement).
Overall, BOSHNAS with the Transformer2vec embeddings performs the best in this
representative design space, outperforming the state-of-the-art (_i.e._ ,
BANANAS on naive embeddings) by 13%.

Figure 4(a) shows the best GLUE score reached by each baseline NAS algorithm
along with the number of models it trained. Again, these runs are performed on
the representative design space described above, using the Transformer2vec
encodings. As can be seen from the figure, BOSHNAS reaches the best GLUE
score. Ablation analysis justifies the need for heteroscedastic modeling and
second-order optimization (see Figure 4(b)). The heteroscedastic model forces
the optimization of the training recipe when the framework approaches optimal
architectural design decisions. Second-order gradients, on the other hand,
help the search avoid local optima and saddle points, and also aid faster
convergence.

Table 3: Comparison between FlexiBERT and baselines. Results are evaluated on
the development set of the GLUE benchmark. We use Matthews correlation for
CoLA, Spearman correlation for STS-B, and accuracy for other tasks. MNLI is
reported on the matched set. Ablation models for BOSHNAS without second-order
gradients (w/o S.) and without using the heteroscedastic model (w/o H.) are
also included. Best (second-best) performance values are in boldface
(underlined). ∗Performance for NAS-BERT10 was not reported on the WNLI dataset
and was reproduced using an equivalent model in our design space. †FlexiBERT-
Mini model that only optimizes performance on the first eight tasks, for fair
comparisons with NAS-BERT.

Model | Parameters | CoLA | MNLI | MRPC | QNLI | QQP | RTE | SST-2 | STS-B | WNLI | Avg.
---|---|---|---|---|---|---|---|---|---|---|---
BERT-Mini (?) | 16.6M | 0 | 74.8 | 71.8 | 84.1 | 66.4 | 57.9 | 85.9 | 73.3 | 62.3 | 64.0
NAS-BERT10 (?) | 10M | 27.8 | 76.0 | 81.5 | 86.3 | 88.4 | 66.6 | 88.6 | 84.8 |  53.7∗ | 72.6
FlexiBERT-Mini (ours, w/o S.) | 7.2M | 16.7 | 72.3 | 72.9 | 81.7 | 76.9 | 64.1 | 80.9 | 77.0 | 65.3 | 67.5
FlexiBERT-Mini (ours, w/o H.) | 20M | 12.3 | 74.4 | 72.3 | 76.4 | 76.3 | 59.5 | 81.2 | 75.4 | 67.8 | 66.2
FlexiBERT-Mini† (ours) | 13.8M | 28.7 | 77.5 | 82.3 | 86.9 | 87.8 | 67.6 | 89.7 | 83.0 | 51.8 | 72.7
FlexiBERT-Mini (ours) | 16.1M | 23.8 | 76.1 | 82.4 | 87.1 | 88.7 | 69.0 | 81.0 | 78.9 | 69.3 | 72.9

Table 3 shows the scores of the ablation models on the GLUE benchmarking
tasks. We refer to the best model obtained from BOSHNAS in the Tiny-to-Mini
space as FlexiBERT-Mini. Once we get the best architecture from the search
process (using the same, albeit limited compute budget for feasible search
times), it is pre-trained and fine-tuned on a larger compute budget (details
in Appendix B.5). As can be seen from the table, FlexiBERT-Mini outperforms
the baseline, NAS-BERT (?), by 0.4% on the GLUE benchmark. Since NAS-BERT
finds the higher-performing architecture while only considering the first
eight GLUE tasks (_i.e._ , without the WNLI dataset), for fair comparisons, we
find a neighboring model in the FlexiBERT design space that only optimizes
performance on the first eight tasks. We call this model FlexiBERT-Mini†. We
see that although FlexiBERT-Mini† does not have the highest GLUE score, it
outperforms NAS-BERT10 by significant margins on the first eights tasks.

Figure 5: Performance frontiers of FlexiBERT on an expanded design space
(under the constraints defined in Table 2) and for traditional homogeneous
models.

Figure 5 demonstrates that FlexiBERT pushes to improve the performance
frontier relative to traditional homogeneous architectures. In other words,
the best-performing models in the expanded (Tiny-to-Mini) space outperform
traditional models for the same number of parameters. Here, the homogeneous
models incorporate the same design decisions for all encoder layers, even with
the expanded set of operations (_i.e._ , including convolutional and LT-based
attention operations). FlexiBERT-Mini has 3% fewer parameters than BERT-Mini
and achieves 8.9% higher GLUE score. FlexiBERT achieves 3% higher performance
than the best homogeneous model while the model with equivalent performance
achieves 2.6$\times$ smaller size.

###  4.4 Best Architecture in the Design Space

After running BOSHNAS for each level of the hierarchy, we get the respective
best-performing models, whose model cards are presented in Appendix B.6. From
these best-performing models, we can extract the following rules that lead to
high-performing transformer architectures:

  * •

Models with DCT in the deeper layers are preferable for higher performance on
the GLUE benchmark.

  * •

Models with more attention heads, but smaller hidden dimension, are preferable
in the deeper layers.

  * •

Feed-forward networks with larger widths, but smaller depth, are preferable in
the deeper layers.

Table 4: Comparison between FlexiBERT-Large (outside of the constraints
defined in Table 2) and baselines on GLUE score. ∗GLUE scores reported do not
consider the WNLI dataset.

Model | Parameters | GLUE score
---|---|---
RoBERTa (?) | 345M | 88.5
FNet-Large (?) | 357M |  81.9∗
AutoTinyBERT (?) | 85M |  81.2∗
DynaBERT (?) | 345M |  81.6∗
NAS-BERT60 (?) | 60M |  83.2∗
AutoBERT-Zero Large (?) | 318M |  84.5∗
FlexiBERT-Large (ours) | 319M | 89.1/90.2∗

Using these guidelines, we extrapolate the model card for FlexiBERT-Mini to
get the design decisions for FlexiBERT-Large, which is an equivalent
counterpart of BERT-Large (?). Appendix B.6 presents the approach for
extrapolation of hyperparameter choices from FlexiBERT-Mini to obtain
FlexiBERT-Large. We train FlexiBERT-Large with the larger compute budget (see
Appendix B.5) and show its GLUE score in Table 4. FlexiBERT-Large outperforms
the baseline RoBERTa by 0.6% on the entire GLUE benchmarking suite, and
AutoBERT-Zero Large by 5.7% when only considering the first eight tasks.

Just like FlexiBERT-Large is the BERT-Large counterpart of FlexiBERT-Mini, we
similarly form the BERT-Small and BERT-Base equivalents (?). Figure 6 presents
the performance frontier of these FlexiBERT models with different baseline
works. As can be seen, FlexiBERT consistently outperforms the baselines for
different constraints on model size, thanks to its search in a vast,
_heterogeneous_ , and _flexible_ design space of architectures.

Figure 6: Performance of FlexiBERT and other baseline methods on various GLUE
tasks: (a) SST-2, (b) QNLI, (c) MNLI (accuracy of MNLI-m is plotted), and (d)
CoLA.

##  5 Conclusion

In this work, we presented FlexiBERT, a suite of _heterogeneous_ and
_flexible_ transformer models. We characterized the effects of this expanded
design space and proposed a novel Transformer2vec embedding scheme to train a
surrogate model that searches the design space for high-performance models. We
described a novel NAS algorithm, BOSHNAS, and showed that it outperforms the
state-of-the-art by 13%. The FlexiBERT-Mini model searched in this design
space has a GLUE score that is 8.9% higher than BERT-Mini, while requiring 3%
lower parameters. It also beats the baseline, NAS-BERT10 by 0.4%. A FlexiBERT
model with equivalent performance as the best homogeneous model achieves
2.6$\times$ smaller size. FlexiBERT-Large outperforms the state-of-the-art
models by at least 5.7% average accuracy on the first eight tasks in the GLUE
benchmark.

## Acknowledgments

This work was supported by NSF Grant No. CNS-1907381 and CCF-2203399. The
experiments reported in this paper were substantially performed on the
computational resources managed and supported by Princeton Research Computing
at Princeton University. We also thank Xiaorun Wu for initial discussions.

##  A Background

Here, we discuss some supplementary background concepts.

###  A.1 Self-Attention

Traditionally, transformers have relied on the SA operation. It is basically a
trainable associative memory. We depict the vanilla SA operation as SDP and
introduce the WMA operation in our design space as well. For a source vector
$\mathbf{s}$ and a hidden-state vector $\mathbf{h}$:

$\text{SDP}:=\frac{\mathbf{s}^{\top}\mathbf{h}}{\sqrt{d}},\
\text{WMA}:=\mathbf{s}^{\top}\mathbf{W}_{a}\mathbf{h}$

where $d$ is the dimension of the source vector and $\mathbf{W}_{a}$ is a
trainable weight matrix in the attention layer. Naturally, a WMA layer is more
expressive than an SDP layer.

The SA mechanism used in the context of transformers also involves the softmax
function and matrix multiplication. More concretely, in a multi-headed SA
operation (with $n$ heads), there are four matrices:
$\mathbf{W}^{q}_{i}\in\mathbb{R}^{d_{inp}\times h/n}$,
$\mathbf{W}^{k}_{i}\in\mathbb{R}^{d_{inp}\times h/n}$,
$\mathbf{W}^{v}_{i}\in\mathbb{R}^{d_{inp}\times h/n}$, and
$\mathbf{W}^{o}_{i}\in\mathbb{R}^{h/n\times d_{out}}$, and it takes the hidden
states of the previous layer as input $\mathbf{H}\in\mathbb{R}^{N_{T}\times
d_{inp}}$, where $i$ refers to an attention head, $d_{inp}$ is the input
dimension, $d_{out}$ is the output dimension, and $h$ is the hidden dimension.
The output of the attention head ($\mathbf{H}_{i}\in\mathbb{R}^{N_{T}\times
d_{out}}$) is then calculated as follows:

$\displaystyle\mathbf{Q}_{i},\mathbf{K}_{i},\mathbf{V}_{i}=\mathbf{H}\mathbf{W}^{q}_{i},\mathbf{H}\mathbf{W}^{k}_{i},\mathbf{H}\mathbf{W}^{v}_{i}$
$\displaystyle\mathbf{H}_{i}=\text{softmax}\left(\frac{\mathbf{Q}_{i}\mathbf{K}_{i}^{\intercal}}{\sqrt{h}}\right)\mathbf{V}_{i}\mathbf{W}^{o}_{i}$

For traditional _homogenous_ transformer models, $d_{out}$ has to equal
$d_{inp}$ (usually, $d_{inp}=d_{out}=h$) due to the residual connections.
However, thanks to the _relative_ and _trained_ positional encodings and the
added projection layer at the end of each encoder
($\mathbf{W}^{p}\in\mathbb{R}^{d_{out}\times d_{p}}$), we can relax this
constraint. This leads to an expansion of the _flexibility_ of transformer
models in the FlexiBERT design space.

###  A.2 Improving BERT’s performance

BERT is one of the most widely used transformer architectures (?). Researchers
have improved BERT’s performance by revamping the pre-training technique.
RoBERTa proposed a more robust pre-training approach to improve BERT’s
performance by considering _dynamic_ masking in the Masked Language Modeling
(MLM) objective (?). Functional improvements have also been proposed for pre-
training – XLNet introduced Permuted Language Modeling (PLM) (?) and MPNet
extended it by unifying MLM and PLM techniques (?). Other approaches,
including denoising autoencoders, have also been proposed (?).

On the other hand, ? (?) consider optimizing the set of _architectural_ design
decisions for BERT – number of encoder layers $l$, size of hidden embeddings
$h$, number of attention heads $a$, size of the hidden layer in the feed-
forward network $f$, etc. However, it is only concerned with pruning BERT and
does not target optimization of accuracy over different tasks. Further, it has
a limited search space consisting of only _homogeneous_ models.

##  B Experimental Details

We present the details of the experiments performed next.

###  B.1 Possible Compute Blocks

Based on the design space presented in Table 2, we consider all possible
compute blocks, as presented next:

  * •

For layer $j$, when the operation is SA, we have two or four heads with:
$h$-128/SA-SDP, $h$-128/SA-WMA, $h$-256/SA-SDP, and $h$-256/SA-WMA. If the
encoder layer has an LT operation, then we have two or four heads with:
$h$-128/LT-DFT, $h$-128/LT-DCT, $h$-256/LT-DFT, and $h$-256/LT-DCT; the latter
entry being the type of LT operation. For a _convolutional_ (DSC) operation,
we have two or four heads with: $h$-128/DSC-5, $h$-128/DSC-9, $h$-256/DSC-5,
and $h$-256/DSC-9; the latter entry referring to the kernel size.

  * •

For layer $j$, the size of the hidden layer in the feed-forward network is
either 512 or 1024. Also, the feed-forward network may either have just one
hidden layer or a stack of three layers. At higher levels of the hierarchy in
the hierarchical search framework (details in Section 3.2), all the layers in
the stack of hidden layers have the same dimension, until we relax this
constraint in the last leg of the hierarchy.

  * •

Other blocks like: Add&Norm, Input, and Output.

###  B.2 Knowledge Transfer

Knowledge transfer has been used in recent works, but restricted to long
short-term memories and simple recurrent neural networks (?). ? (?) train a
super-transformer and share its weights with smaller models. However, this is
not feasible for diverse _heterogeneous_ and _flexible_ architectures. We
propose the use of knowledge transfer in transformers for the first time, to
the best of our knowledge, by comparing weights with computational graphs of
_nearby_ models. Furthermore, previous works only consider a static training
recipe for all the models in the design space, an assumption we relax in our
experiments. We _directly_ fine-tune models for which _nearby_ models are
already pre-trained. We test for this using the _biased overlap_ metric
defined in Section 3.3. Figure 7 presents the time gains from knowledge
transfer. Since some percentage of models could directly be fine-tuned, thanks
to their neighboring pre-trained models, we were able to speed up the overall
training time by 38%.

Figure 7: Bar plot showing average time for training a transformer model (in
GPU-hours) with and without knowledge transfer. (a) Pre-train + Fine-tune:
total training time. (b) Direct Fine-tune: training time for a pre-trained
model. (c) Knowledge Transfer: training using weight transfer from a trained
_nearby_ model, gives 38% speedup. Plotted with 90% confidence intervals.
Figure 8: Crossover between two parent models yields a finer-grained design
space. Each stack configuration in the children is derived from the product of
the parent design choices at the same depth.

###  B.3 Crossover between Transformer Models

We obtain new transformer models of the subsequent level in the hierarchy by
taking a crossover between the best models in the previous level (which had
layers per stack $=s$) and their neighbors. The stack configuration of the
children is chosen from all unique hyperparameter values present in the parent
models at the same depth. We present a simple example of this scheme in Figure
8. The design space of permissible operation blocks for layers in the stack,
$s$, is computed by the product of the respective design choices of the
parents for that stack. These layers are then independently formed with the
new constraint of $s/2$ layers having the same choice of hyperparameter
values. Expanding the design space in such a fashion retains the original
hyperparameters that give good performance while also exploring the internal
representations learned by combinations of the hyperparameters at the same
level.

###  B.4 BOSHNAS Training Flow

Different surrogate models in the BOSHNAS pipeline ($f_{S},g_{S}$, and
$h_{S}$) have been presented in the order of flow in Figure 9. As explained in
Section 3.4, the NPN network ($f_{S}$) models the model performance and the
_aleatoric_ uncertainty, and the student network ($h_{S}$) models the
_epistemic_ uncertainty from the teacher network ($g_{S}$).

Figure 9: Overview of the BOSHNAS pipeline. Variables have been defined in
Section 3.4

###  B.5 Model Training

We pre-train our models with a combination of publicly available text corpora,
viz. BookCorpus (BookC) (?), Wikipedia English (Wiki), OpenWebText (OWT) (?),
and CC-News (CCN) (?). Most training hyperparameters are borrowed from
RoBERTa. We set the batch size to 256, learning rate warmed up over the first
$10,000$ steps to its peak value at $1\times 10^{-5}$ that then decays
linearly, weight decay to $0.01$, Adam scheduler’s parameters $\beta_{1}=0.9$,
$\beta_{2}=0.98$ (shown to improve stability; ?), $\epsilon=1\times 10^{-6}$,
and run pre-training for $1,000,000$ steps.

Once the best models are found, we pre-train and fine-tune the selected models
with a larger compute budget. For pre-training, we add the C4 dataset (?) and
train for $3,000,000$ steps before fine-tuning. We also fine-tune on each GLUE
task for $10$ epochs instead of $5$ (further details are given below). This
was done for the FlexiBERT-Mini and FlexiBERT-Large models. Table 5 shows the
improvement in performance of FlexiBERT-Mini that was trained using knowledge
transfer (where the weights were transferred from a nearby trained model)
after additional training. When compared to the model directly fine-tuned
after knowledge transfer, we see only a marginal improvement when we pre-train
from scratch. This reaffirms the advantage of knowledge transfer, that it
reduces training time (see Appendix B.2) with negligible loss in performance.
Training with a larger compute budget further improves performance on the GLUE
benchmark, validating the importance of data size and diversity in pre-
training (?). Running a full-fledged BOSHNAS on the larger design space
(_i.e._ , with layers from 2 to 24, Tiny-to-Large) can be an easy extension of
this work.

While running BOSHNAS, we fine-tune our models on the nine GLUE tasks over
five epochs and a batch size of $64$ where early stopping is implemented. We
also run automatic hyperparameter tuning for the fine-tuning process using the
Tree-structured Parzen Estimator algorithm (?). The learning rate is randomly
selected logarithmically in the [$2\times 10^{-5},\;5\times 10^{-4}$] range,
and the batch size in $\\{32,64,128\\}$ uniformly. Table 6 shows the best
hyperparameters for fine-tuning of each GLUE task selected using this auto-
tuning technique. This hyperparameter optimization uses random initialization
every time, which results in variation in performance each time the model is
queried (see _aleatoric_ uncertainty explained in Section 3.4).

Table 5: Performance of FlexiBERT-Mini from BOSHNAS after knowledge transfer
from a nearby trained model, and after pre-training from scratch along with a
larger compute budget.

Model | Pre-training data | Pre-training steps | Fine-tuning epochs | GLUE score
---|---|---|---|---
FlexiBERT-Mini |  |  |  |
w/ knowledge transfer | BookC, Wiki, OWT, CCN | 1,000,000 | 5 | 69.7
\+ pre-training from scratch | BookC, Wiki, OWT, CCN | 1,000,000 | 5 | 70.4
\+ larger compute budget | BookC, Wiki, OWT, CCN, C4 | 3,000,000 | 10 | 72.9

Table 6: Hyperparameters used for fine-tuning FlexiBERT-Mini on the GLUE tasks. Task | Learning rate | Batch size
---|---|---
CoLA | $2.0\times 10^{-4}$ | 64
MNLI | $9.4\times 10^{-5}$ | 64
MRPC | $2.23\times 10^{-5}$ | 32
QNLI | $5.03\times 10^{-5}$ | 128
QQP | $3.7\times 10^{-4}$ | 64
RTE | $1.9\times 10^{-4}$ | 128
SST-2 | $1.2\times 10^{-4}$ | 128
STS-B | $7.0\times 10^{-5}$ | 32
WNLI | $4.0\times 10^{-5}$ | 128

We have included baselines trained with the _pre-training_ \+ _fine-tuning_
procedure as proposed by ? (?) for like-for-like comparisons, and not the
_knowledge distillation_ counterparts (?). Nevertheless, FlexiBERT is
orthogonal to (and thus can easily be combined with) knowledge distillation
because FlexiBERT focuses on searching the best architecture while knowledge
distillation focuses on better training of a given architecture.

All models were trained on NVIDIA A100 GPUs and 2.6 GHz AMD EPYC Rome
processors. The entire process of running BOSHNAS for all levels of the
hierarchy took around 300 GPU-days of training.

###  B.6 Best-performing Models

From different hierarchy levels ($s=2,1$, and $1^{*}$), we get the respective
best-performing models after running BOSHNAS as follows:

$\displaystyle
s=2:\Big{\\{}l:4,o:[\text{LT},\text{LT},\text{LT},\text{LT}],h:[256,256,256,256],n:[4,4,2,2],$
$\displaystyle
f:[[1024],[1024],[512,512,512],[512,512,512]],p:[\text{DCT},\text{DCT},\text{DCT},\text{DCT}]\Big{\\}}$
$\displaystyle
s=1:\Big{\\{}l:4,o:[\text{SA},\text{SA},\text{LT},\text{LT}],h:[256,256,128,128],n:[2,2,4,4],$
$\displaystyle
f:[[512,512,512],[512,512,512],[1024],[1024]],p:[\text{SDP},\text{SDP},\text{DCT},\text{DCT}]\Big{\\}}$

where, in the last leg of the hierarchy, the stack length is $1$, but the
feed-forward stacks are also heterogeneous (see Section 3.2). Both $s=1$ and
$s=1^{*}$ gave the same solution despite finer granularity in the latter case.
Thus, the second model card above is that of FlexiBERT-Mini.

The model cards of the FlexiBERT-Mini ablation models, as presented in Table
3, are given below:

$\displaystyle\text{(w/o
S.)}:\Big{\\{}l:2,o:[\text{SA},\text{SA}],h:[128,128],n:[4,4],f:[[1024],[1024]],p:[\text{SDP},\text{WMA}]\Big{\\}}$
$\displaystyle\text{(w/o
H.)}:\Big{\\{}l:4,o:[\text{LT},\text{LT},\text{SA},\text{SA}],h:[256,256,128,128],n:[4,4,4,4],$
$\displaystyle
f:[[1024,1024,1024],[1024,1024,1024],[512,512,512],[512,512,512]],$
$\displaystyle p:[\text{DCT},\text{DCT},\text{SDP},\text{SDP}]\Big{\\}}$

Figure 10 shows a working schematic of the design choices in the FlexiBERT-
Mini and FlexiBERT-Large models. As explained in Section 4.4, FlexiBERT-Large
was formed by extrapolating the design choices in FlexiBERT-Mini to obtain a
BERT-Large counterpart (?).

Figure 10: Obtained FlexiBERT models after running the BOSHNAS pipeline: (a)
FlexiBERT-Mini, and its design choices extrapolated to obtain (b) FlexiBERT-
Large.

## References

  * Abu-Aisheh et al. Abu-Aisheh, Z., Raveaux, R., Ramel, J.-Y., and Martineau, P. (2015).  An exact graph edit distance algorithm for solving pattern recognition problems.  In Proceedings of the International Conference on Pattern Recognition Applications and Methods, Vol. 1, pp. 271–278.
  * Akiba et al. Akiba, T., Sano, S., Yanase, T., Ohta, T., and Koyama, M. (2019).  Optuna: A next-generation hyperparameter optimization framework.  In Proceedings of the 25th ACM SIGKDD International Conference on Knowledge Discovery & Data Mining, pp. 2623–2631.
  * Chen et al. Chen, D., Li, Y., Qiu, M., Wang, Z., Li, B., Ding, B., Deng, H., Huang, J., Lin, W., and Zhou, J. (2021).  AdaBERT: Task-adaptive BERT compression with differentiable neural architecture search.  In Proceedings of the 29th International Joint Conference on Artificial Intelligence, pp. 2463–2469.
  * Cheng et al. Cheng, H.-P., Zhang, T., Zhang, Y., Li, S., Liang, F., Yan, F., Li, M., Chandra, V., Li, H., and Chen, Y. (2021).  NASGEM: Neural architecture search via graph embedding method.  In Proceedings of the AAAI Conference on Artificial Intelligence, Vol. 35, pp. 7090–7098.
  * Conneau and Lample Conneau, A., and Lample, G. (2019).  Cross-lingual language model pretraining.  In Advances in Neural Information Processing Systems, Vol. 32, pp. 7059–7069.
  * Dang et al. Dang, D.-C., Eremeev, A., and Lehre, P. K. (2021).  Escaping local optima with non-elitist evolutionary algorithms.  In Proceedings of the AAAI Conference on Artificial Intelligence, Vol. 35, pp. 12275–12283.
  * Devlin et al. Devlin, J., Chang, M.-W., Lee, K., and Toutanova, K. (2019).  BERT: Pre-training of deep bidirectional transformers for language understanding.  In Proceedings of the Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Vol. 1, pp. 4171–4186.
  * Gal and Ghahramani Gal, Y., and Ghahramani, Z. (2016).  Dropout as a Bayesian approximation: Representing model uncertainty in deep learning.  In Proceedings of The 33rd International Conference on Machine Learning, Vol. 48, pp. 1050–1059.
  * Gao et al. Gao, J., Xu, H., Shi, H., Ren, X., Yu, P. L. H., Liang, X., Jiang, X., and Li, Z. (2021).  AutoBERT-Zero: Evolving BERT backbone from scratch.  CoRR, abs/2107.07445.
  * Gokaslan and Cohen Gokaslan, A., and Cohen, V. (2019).  OpenWebText corpus.  http://Skylion007.github.io/OpenWebTextCorpus.
  * He et al. He, X., Zhao, K., and Chu, X. (2021).  AutoML: A survey of the state-of-the-art.  Knowledge-Based Systems, 212, 106622.
  * Hou et al. Hou, L., Huang, Z., Shang, L., Jiang, X., Chen, X., and Liu, Q. (2020).  DynaBERT: Dynamic BERT with adaptive width and depth.  In Advances in Neural Information Processing Systems, Vol. 33, pp. 9782–9793.
  * Huang et al. Huang, C.-Z. A., Vaswani, A., Uszkoreit, J., Simon, I., Hawthorne, C., Shazeer, N., Dai, A. M., Hoffman, M. D., Dinculescu, M., and Eck, D. (2018).  Music transformer: Generating music with long-term structure.  In Proceedings of the International Conference on Learning Representations.
  * Jiang et al. Jiang, Z.-H., Yu, W., Zhou, D., Chen, Y., Feng, J., and Yan, S. (2020).  ConvBERT: Improving BERT with span-based dynamic convolution.  In Advances in Neural Information Processing Systems, Vol. 33, pp. 12837–12848.
  * Khetan and Karnin Khetan, A., and Karnin, Z. (2020).  schuBERT: Optimizing elements of BERT.  In Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics, pp. 2807–2818.
  * Kruskal Kruskal, J. B. (1964).  Multidimensional scaling by optimizing goodness of fit to a nonmetric hypothesis.  Psychometrika, 29(1), 1–27.
  * Lee-Thorp et al. Lee-Thorp, J., Ainslie, J., Eckstein, I., and Ontanon, S. (2021).  FNet: Mixing tokens with Fourier transforms.  CoRR, abs/2105.03824.
  * Lewis et al. Lewis, M., Liu, Y., Goyal, N., Ghazvininejad, M., Mohamed, A., Levy, O., Stoyanov, V., and Zettlemoyer, L. (2020).  BART: Denoising sequence-to-sequence pre-training for natural language generation, translation, and comprehension.  In Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics, pp. 7871–7880.
  * Liu et al. Liu, Y., Ott, M., Goyal, N., Du, J., Joshi, M., Chen, D., Levy, O., Lewis, M., Zettlemoyer, L., and Stoyanov, V. (2019).  RoBERTa: A robustly optimized BERT pretraining approach.  CoRR, abs/1907.11692.
  * Lu et al. Lu, Z., Whalen, I., Boddeti, V., Dhebar, Y., Deb, K., Goodman, E., and Banzhaf, W. (2019).  NSGA-Net: Neural architecture search using multi-objective genetic algorithm.  In Proceedings of the Genetic and Evolutionary Computation Conference, pp. 419–427.
  * Luong et al. Luong, T., Pham, H., and Manning, C. D. (2015).  Effective approaches to attention-based neural machine translation.  In Proceedings of the Conference on Empirical Methods in Natural Language Processing, pp. 1412–1421.
  * Mackenzie et al. Mackenzie, J., Benham, R., Petri, M., Trippas, J. R., Culpepper, J. S., and Moffat, A. (2020).  CC-News-En: A large English news corpus.  In Proceedings of the 29th ACM International Conference on Information & Knowledge Management, pp. 3077–3084.
  * Mazzawi et al. Mazzawi, H., Gonzalvo, X., Kracun, A., Sridhar, P., Subrahmanya, N., Lopez-Moreno, I., Park, H., and Violette, P. (2019).  Improving keyword spotting and language identification via neural architecture search at scale.  In Proceedings of Interspeech, pp. 1278–1282.
  * Mikolov et al. Mikolov, T., Sutskever, I., Chen, K., Corrado, G. S., and Dean, J. (2013).  Distributed representations of words and phrases and their compositionality.  In Advances in Neural Information Processing Systems, Vol. 26, pp. 3111–3119.
  * Narayanan et al. Narayanan, A., Chandramohan, M., Venkatesan, R., Chen, L., Liu, Y., and Jaiswal, S. (2017).  graph2vec: Learning distributed representations of graphs.  CoRR, abs/1707.05005.
  * Pennington et al. Pennington, J., Socher, R., and Manning, C. (2014).  GloVe: Global vectors for word representation.  In Proceedings of the Conference on Empirical Methods in Natural Language Processing, pp. 1532–1543.
  * Pham et al. Pham, H., Guan, M., Zoph, B., Le, Q., and Dean, J. (2018).  Efficient neural architecture search via parameters sharing.  In Proccedings of the 35th International Conference on Machine Learning, pp. 4095–4104.
  * Raffel et al. Raffel, C., Shazeer, N., Roberts, A., Lee, K., Narang, S., Matena, M., Zhou, Y., Li, W., and Liu, P. J. (2019).  Exploring the limits of transfer learning with a unified text-to-text transformer.  CoRR, abs/1910.10683.
  * Real et al. Real, E., Aggarwal, A., Huang, Y., and Le, Q. V. (2019).  Regularized evolution for image classifier architecture search.  In Proceedings of the AAAI Conference on Artificial Intelligence, Vol. 33, pp. 4780–4789.
  * Ru et al. Ru, R., Esperanca, P., and Carlucci, F. M. (2020).  Neural architecture generator optimization.  In Advances in Neural Information Processing Systems, Vol. 33, pp. 12057–12069.
  * Russell and Norvig Russell, S., and Norvig, P. (2010).  Artificial Intelligence: A Modern Approach (3rd edition).  Prentice Hall.
  * Shaw et al. Shaw, P., Uszkoreit, J., and Vaswani, A. (2018).  Self-attention with relative position representations.  In Proceedings of the Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Vol. 2, pp. 464–468.
  * Shervashidze et al. Shervashidze, N., Schweitzer, P., van Leeuwen, E. J., Mehlhorn, K., and Borgwardt, K. M. (2011).  Weisfeiler-Lehman graph kernels.  Journal of Machine Learning Research, 12(77), 2539–2561.
  * Siems et al. Siems, J., Zimmer, L., Zela, A., Lukasik, J., Keuper, M., and Hutter, F. (2020).  NAS-Bench-301 and the case for surrogate benchmarks for neural architecture search.  CoRR, abs/2008.09777.
  * Snoek et al. Snoek, J., Larochelle, H., and Adams, R. P. (2012).  Practical Bayesian optimization of machine learning algorithms.  In Advances in Neural Information Processing Systems, Vol. 25, pp. 2951–2959.
  * So et al. So, D., Mańke, W., Liu, H., Dai, Z., Shazeer, N., and Le, Q. V. (2021).  Searching for efficient transformers for language modeling.  In Advances in Neural Information Processing Systems, Vol. 34, pp. 6010–6022.
  * So et al. So, D. R., Liang, C., and Le, Q. V. (2019).  The evolved transformer.  CoRR, abs/1901.11117.
  * Song et al. Song, K., Tan, X., Qin, T., Lu, J., and Liu, T.-Y. (2020).  MPNet: Masked and permuted pre-training for language understanding.  In Advances in Neural Information Processing Systems, Vol. 33, pp. 16857–16867.
  * Sun et al. Sun, Z., Yu, H., Song, X., Liu, R., Yang, Y., and Zhou, D. (2020).  MobileBERT: A compact task-agnostic BERT for resource-limited devices.  In Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics, pp. 2158–2170.
  * Tan and Le Tan, M., and Le, Q. (2019).  EfficientNet: Rethinking model scaling for convolutional neural networks.  In Proceedings of the 36th International Conference on Machine Learning, pp. 6105–6114.
  * Tuli et al. Tuli, S., Poojara, S. R., Srirama, S. N., Casale, G., and Jennings, N. R. (2021).  COSCO: Container orchestration using co-simulation and gradient based optimization for fog computing environments.  IEEE Transactions on Parallel and Distributed Systems, 33(1), 101–116.
  * Turc et al. Turc, I., Chang, M., Lee, K., and Toutanova, K. (2019).  Well-read students learn better: The impact of student initialization on knowledge distillation.  CoRR, abs/1908.08962.
  * Vaswani et al. Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, Ł., and Polosukhin, I. (2017).  Attention is all you need.  In Advances in Neural Information Processing Systems, Vol. 30, pp. 5998–6008.
  * Wang et al. Wang, A., Singh, A., Michael, J., Hill, F., Levy, O., and Bowman, S. (2018).  GLUE: A multi-task benchmark and analysis platform for natural language understanding.  In Proceedings of the EMNLP Workshop BlackboxNLP: Analyzing and Interpreting Neural Networks for NLP, pp. 353–355.
  * Wang et al. Wang, H., Wu, Z., Liu, Z., Cai, H., Zhu, L., Gan, C., and Han, S. (2020).  HAT: Hardware-aware transformers for efficient natural language processing.  In Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics, pp. 7675–7688.
  * Wang et al. Wang, H., Shi, X., and Yeung, D.-Y. (2016).  Natural-parameter networks: A class of probabilistic neural networks.  In Advances in Neural Information Processing Systems, Vol. 29, pp. 118–126.
  * White et al. White, C., Neiswanger, W., and Savani, Y. (2021a).  BANANAS: Bayesian optimization with neural architectures for neural architecture search.  In Proceedings of the AAAI Conference on Artificial Intelligence, Vol. 35, pp. 10293–10301.
  * White et al. White, C., Zela, A., Ru, B., Liu, Y., and Hutter, F. (2021b).  How powerful are performance predictors in neural architecture search?.  CoRR, abs/2104.01177.
  * Williams Williams, R. J. (1992).  Simple statistical gradient-following algorithms for connectionist reinforcement learning.  Machine Learning, 8(3–4), 229–256.
  * Xu et al. Xu, J., Tan, X., Luo, R., Song, K., Li, J., Qin, T., and Liu, T.-Y. (2021).  NAS-BERT: Task-agnostic and adaptive-size BERT compression with neural architecture search.  In Proceedings of the 27th ACM SIGKDD Conference on Knowledge Discovery & Data Mining, pp. 1933–1943.
  * Xu et al. Xu, K., Hu, W., Leskovec, J., and Jegelka, S. (2019).  How powerful are graph neural networks?.  In Proceedings of the International Conference on Learning Representations.
  * Yang et al. Yang, Z., Dai, Z., Yang, Y., Carbonell, J. G., Salakhutdinov, R., and Le, Q. V. (2019).  XLNet: Generalized autoregressive pretraining for language understanding.  CoRR, abs/1906.08237.
  * Yao et al. Yao, Z., Gholami, A., Shen, S., Mustafa, M., Keutzer, K., and Mahoney, M. (2021).  ADAHESSIAN: An adaptive second order optimizer for machine learning.  In Proceedings of the AAAI Conference on Artificial Intelligence, Vol. 35, pp. 10665–10673.
  * Yin et al. Yin, Y., Chen, C., Shang, L., Jiang, X., Chen, X., and Liu, Q. (2021).  AutoTinyBERT: Automatic hyper-parameter optimization for efficient pre-trained language models.  In Proceedings of the 59th Annual Meeting of the Association for Computational Linguistics and the 11th International Joint Conference on Natural Language Processing, Vol. 1, pp. 5146–5157.
  * Ying et al. Ying, C., Klein, A., Christiansen, E., Real, E., Murphy, K., and Hutter, F. (2019).  NAS-Bench-101: Towards reproducible neural architecture search.  In Proceedings of the 36th International Conference on Machine Learning, Vol. 97, pp. 7105–7114.
  * Yu and Jha Yu, Y., and Jha, N. K. (2022).  SPRING: A sparsity-aware reduced-precision monolithic 3D CNN accelerator architecture for training and inference.  IEEE Transactions on Emerging Topics in Computing, 10(1), 237–249.
  * Zhang et al. Zhang, X., Zhou, X., Lin, M., and Sun, J. (2018).  ShuffleNet: An extremely efficient convolutional neural network for mobile devices.  In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pp. 6848–6856.
  * Zhu et al. Zhu, Y., Kiros, R., Zemel, R., Salakhutdinov, R., Urtasun, R., Torralba, A., and Fidler, S. (2015).  Aligning books and movies: Towards story-like visual explanations by watching movies and reading books.  In Proceedings of the IEEE International Conference on Computer Vision, pp. 19–27.
  * Zoph and Le Zoph, B., and Le, Q. V. (2016).  Neural architecture search with reinforcement learning.  CoRR, abs/1611.01578.
  * Zoph et al. Zoph, B., Vasudevan, V., Shlens, J., and Le, Q. V. (2017).  Learning transferable architectures for scalable image recognition.  CoRR, abs/1707.07012.