LiteTransformerSearch: Training-free On-device Search for Efficient Autoregressive Language Models
Anonymous Author(s)
Abstract
The Transformer architecture is ubiquitously used as the building block of largescale autoregressive language models.
However, finding architectures with theoptimal trade-off between task performance (perplexity) and hardware constraintslike peak memory utilization and latency is non-trivial.
This is exacerbated by theproliferation of various hardware.
We leverage the somewhat surprising empiricalobservation that the number of decoder parameters in autoregressive Transformershas a high rank correlation with task performance, irrespective of the architecturetopology.
This observation organically induces a simple search algorithm that usesdecoder parameters as a proxy for perplexity without need for any model training.
The search phase of our training-free algorithm, dubbed Lightweight Transformer
Search (LTS), can be run directly on target devices since it does not require GPUs.
We focus on extracting the pareto-frontier of perplexity versus any hardware performance cost such as latency and/or memory, using on-target-device measurements.
We evaluate LTS on diverse devices from ARM CPUs to NVIDIA GPUs andtwo popular autoregressive Transformer backbones:
GPT-2 and Transformer-XL.
Results show that the perplexity of 16 -layer GPT-2 and Transformer-XL can beachieved with up to 1 .
6 × , 2 . 5 × faster runtime and 1 . 3 × , 2 . {+
 ×+} lower peak memoryutilization.
LTS extracts the pareto-frontier in under 3 hours, running on a commodity laptop.
We effectively remove the carbon footprint of training during search forhundreds of GPU hours, offering a strong simple baseline for future NAS methodsin autoregressive language modeling.
Introduction
The Transformer architecture [29] has been used as the de-facto building block of most pre-trainedlanguage models like GPT [3].
A common problem arises when one tries to create smaller versionsof Transformer models for edge or real-time applications (e.g. text prediction) with strict memoryand latency constraints: it is not clear what the architectural hyperparameters should be, e.g., numberof attention heads, number of layers, embedding dimension, and the inner dimension of the feedforward network, etc.
This problem is exacerbated if each individual Transformer layer is allowedthe freedom to have different values for these settings.
This results in a combinatorial explosion ofarchitectural hyperparameter choices and a large heterogeneous search space.
For instance, the searchspace considered in this paper consists of over 10 54 possible architectures.
Neural Architecture Search (NAS) is an organic solution due to its ability to automatically searchthrough candidate models with multiple conflicting objectives like latency vs. task performance.
The central challenge in NAS is prohibitively expensive function evaluation, i.e., evaluating eacharchitecture requires training it on the dataset at hand.
Thus it is often infeasible to evaluate morethan a handful of architectures during the search phase.
Supernets [21] have emerged as a dominantparadigm in NAS which combine all possible architectures into a single graph and jointly train themusing weight-sharing.
Nevertheless, supernet training imposes constraints on the expressiveness ofthe search space [20] and is often memory-hungry [37, 4, 36] as it creates large networks duringsearch.
Additionally, training supernets is non-trivial as children architectures may interfere witheach other and rank correlation between sub-architectures is not preserved [20] 1 .
[Figure]
We consider a different approach by proposing a training-free proxy that provides a highly accurateranking of candidate architectures during NAS without need for costly function evaluation or supernets.
Our scope is NAS for efficient Transformer-based autoregressive language models.
We design alightweight search method which is target hardware-aware and outputs a gallery of models on thepareto-frontier of perplexity versus hardware metrics.
We term this method Lightweight Transformer
Search (LTS).
LTS relies on our somewhat surprising observation: the decoder parameter count hasa high rank correlation with the perplexity of fully trained autoregressive Transformers.
Given a set of Transformer-based autoregressive language models, one can accurately rank themusing decoder parameter count as the proxy for perplexity.
Our observations are also well-alignedwith the power laws in [14], shown for homogeneous autoregressive Transformers, i.e., when alldecoder layers have the same configuration.
We provide rigorous experiments to empirically establishthe high rank correlation between perplexity and decoder parameter count for both homogeneous andheterogeneous search spaces.
The above phenomenon coupled with the fact that a candidate architecture’s hardware performancecan be measured on the target device 2 leads to a training-free search procedure: pick one’s favoritediscrete search algorithm (e.g. evolutionary search), sample candidate architectures from the searchspace; count their decoder parameters as a proxy for task performance (i.e., perplexity); measure theirhardware performance (e.g., latency and memory) directly on the target device; and progressivelycreate a pareto-frontier estimate.
While we have chosen a reasonable search algorithm in this work,one can plug-and-play any pareto-frontier search method such as those in [12].
Building upon these insights, Figure 1 shows a high-level overview of LTS.
We design the firsttraining-free Transformer search that is performed entirely on the target (constrained) platform.
Assuch, LTS easily performs a multi-objective NAS where several underlying hardware performancemetrics, e.g., latency and peak memory utilization, are simultaneously optimized.
Using our trainingfree proxy, we extract the 3 -dimensional pareto-frontier of perplexity versus latency and memory in arecord breaking time of < 3 hours on a commodity Intel Core i7 CPU.
Notably, LTS eliminates thecarbon foot print from hundreds of GPU hours of training associated with legacy NAS methods.
To corroborate the effectiveness of our proxy, we train over 2800 Transformer architectures ontwo large language modeling benchmark datasets, namely, WikiText-
[19] and One Billion
Word [5].
We use LTS to search for pareto-optimal architectural hyperparameters in two popularlyused autoregressive Transformer backbones, namely, Transformer-XL [6] and GPT-2 [22].
We believedecoder parameter count should be regarded as a competitive baseline for evaluating Transformer
NAS, both in terms of ranking capabilities and easy computation.
We will open-source our code andour extensive library of trained models to foster future NAS research on Transformers.
Related Work
Here, we discuss literature on automated search for Transformer architectures in the language domain.
We refer to extensive surveys on NAS [9, 34] for a broader overview of the field.
Decoder-only Architectures.
So et al.
[24] search over TensorFlow programs that implement anautoregressive language model via evolutionary search.
Since most random sequences of programseither have errors or underperform, the search has to be seeded with the regular Transformer architecture, termed “Primer”.
As opposed to “Primer” which uses large computation to search a generalspace, our aim is to efficiently search the “backbone” of traditional decoder-only Transformers.
Additionally, the objective in “Primer” is to find models that train faster.
Our objective for NAS, however,is to deliver pareto-frontiers for inference, with respect to perplexity and hardware constraints.
Encoder-only Architectures.
Relative to decoder-only autoregressive language models, encoderonly architectures like BERT [7] have received much more recent attention from the NAS community.
NAS-BERT [35] trains a supernet to efficiently search for masked language models (MLMs) whichare compressed versions of the standard BERT, Such models can then be used in downstream tasks asis standard practice.
Similar to NAS-BERT, [36] also train a supernet to conduct architecture searchwith the aim of finding more efficient BERT variants.
They find interesting empirical insights intosupernet training issues like differing gradients at the same node from different child architecturesand different tensors as input and output at every node in the supernet.
The authors propose fixes thatdrastically improves supernet training.
Tsai et al.
[28], Yin et al.
[38], Gao et al.
[10] also conductvariants of supernet training with the aim of finding more efficient BERT models.
Encoder-Decoder Related:
Applying the well-known DARTS [16] approach to Transformer searchspaces leads to memory-hungry supernets.
To mitigate this issue, [39] propose a multi-split reversiblenetwork and a memory efficient backpropagation algorithm.
One of the earliest papers that applieddiscrete NAS to Transformer search spaces was by So et al.
[23], who use a modified form ofevolutionary search.
Due to the expense of directly performing discrete search on the search space,this work incurs extremely large computation overhead.
Follow-up work by [32] uses the Once-ForAll [4] approach to train a supernet for encoder-decoder architectures used in machine translation.
Search is performed on subsamples of the supernet that inherit weights to estimate task accuracy.
Foreach target device, the authors train a small neural network regressor on thousands of architectures toestimate latency.
As opposed to using a latency estimator, LTS evaluates the latency of each candidatearchitecture on the target hardware.
Notably, by performing the search directly on the target platform,
LTS can easily incorporate various hardware performance metrics, e.g., peak memory utilization, forwhich accurate estimators may not exist.
To the best of our knowledge, such holistic integration ofmultiple hardware metrics in Transformer NAS has not been explored previously.
Lightweight Transformer Search
We perform an evolutionary search over candidate architectures to extract models that lie on thepareto-frontier.
In contrast to the vast majority of prior methods that deliver a single architecture fromthe search space, our search is performed over the entire pareto, generating architectures with a widerange of latency, peak memory utilization, and perplexity with one round of search.
This alleviatesthe need to repeat the NAS algorithm for each hardware performance constraint.
To evaluate candidate models during the search, LTS uses a training-free proxy for the validationperplexity.
By incorporating training-free evaluation metrics, LTS, for the first time, performsthe entire search directly on the target (constrained) hardware.
Therefore, we can use real measurements of hardware performance during the search.
Algorithm 1 outlines the iterative process
Algorithm
Input: Search space D , n iter Output: Perplexity-latency-memory pareto-frontier F 1 L , M , P , F ← ∅ , ∅ , ∅ , ∅ 2 while N ≤ n iter do
F ′ ← Subsample( F )
S N ← EA ( F ′ , D ) // hardware profiling
L ← L (cid:83) Latency( S N )
M ← M (cid:83) Memory( S N ) // estimate perplexity
P ← P (cid:83) Proxy( S N ) // update the pareto front
F ← LowerConvexHull( P , L , M ) performed in LTS 3 for finding candidate architectures in the search space ( D ), that lieon the 3 -dimensional pareto-frontier of perplexity versus latency and memory ( F ).
Ateach iteration, a set of points ( F ′ ) are subsampled from the current pareto-frontier.
Anew batch of architectures ( S N ) are thensampled from F ′ using evolutionary algorithms ( EA ( . ) ).
The new samples are evaluated in terms of latency ( L ), peak memoryutilization ( M ), and validation perplexity( P ).
Latency and memory are measureddirectly on the target hardware while theperplexity is indirectly estimated using ouraccurate and training-free proxy methods.
Finally, the pareto-frontier is recalibrated using the lower convex hull of all sampled architectures.
In the context of multi-objective NAS, pareto-frontier points are those where no single metric (e.g.,perplexity, latency, and memory) can be improved without degrading at least one other metric [12].
To satisfy application-specific needs, optional upper bounds can be placed on the latency and/ormemory of sampled architectures during search.
Search Space.
Figure 2 shows all elastic parameters in LTS search space, namely, numberof layers (n layer ), number of attention heads (n head ), decoder output dimension (d model ), innerdimension of the feed forward network (d inner ), embedding dimension (d embed ), and the division factor ( k ) of adaptive embedding [2].
These architectural parameters are compatible withpopularly used autoregressive Transformer backbones 4 , e.g., GPT.
We adopt a heterogeneous search space where the backbone parameters aredecided on a per-layer basis.
This is in contrastto the homogeneous structure commonly usedin Transformers [6, 3], which reuses the sameconfiguration for all layers.
Compared to homogeneous models, the flexibility associated withheterogeneous architectures enables them to obtain much better hardware performance underthe same perplexity budget (see Section 4.4).
Heterogeneous search space was previously explored in [32].
However, due to the underlying supernet structure, not all design parameters canchange freely.
As an example, the dimensionality of the Q, K, V vectors inside the encoder anddecoder layers is fixed to a large value of 512 to accommodate inheritance from the supernet.
Oursearch space, however, allows exploration of all internal dimensions without constraints.
By notrelying on the supernet structure, our search space easily encapsulates various Transformer backboneswith different configurations of the input/output embedding layers and elastic internal dimensions.
LTS searches over the following values for the architectural parameters in our backbones: n layer ∈{ 2 , . .
, 16 | 1 } 5 , d model ∈ { 128 , . .
. , 1024 | 64 } , d inner ∈ { 128 , . .
4096 | 64 } , and n head ∈ { 2 , 4 , 8 } .
Additionally we explore adaptive input embedding [2] with d embed ∈ { 128 , 256 , 512 } and factork ∈ { 1 , 2 , 4 } .
Once a d model is sampled, we adjust the lower bound of the above range for d innerto 2 × d model .
Encoding this heuristic inside the search ensures that the acquired models will notsuffer from training collapse.
Our heterogeneous search space encapsulates more than 10 54 differentarchitectures.
Such high dimensionality further validates the critical need for training-free NAS.
3.1 Training-free Architecture Ranking
▶ Low-cost Ranking Proxies.
Recently, Abdelfattah et al.
[1] utilize the summation of pruningscores over all model weights as the ranking proxy for Convolutional Neural Networks (CNNs),where a higher score corresponds to higher architecture rank in the search space.
White et al.
[33]analyze these and more recent proxies and find that no particular proxy performs consistently wellover various tasks and baselines, while number of parameters and flops proxies are quite competitive.
However they did not include Transformer-based search spaces in their analysis.
To the best ofour knowledge, low-cost (pruning-based) proxies have not been evaluated on Transformer searchspaces in the language domain.
Note that one cannot naively apply these proxies to language models.
Specifically, since the embedding layer in Transformers is equivalent to a lookup operation, specialcare must be taken to omit this layer from the proxy computation.
Using this insight, we perform thefirst systematic study of low-cost proxies for NAS on autoregressive Transformers for text prediction.
We leverage various pruning metrics, namely, grad_norm , snip [15], grasp [31], fisher [27],and synflow [25].
We also study jacob_cov [18] and relu_log_det [17] which are low-costscoring mechanisms proposed for NAS on CNNs in vision tasks.
While these low-cost techniquesdo not perform model training, they require forward and backward passes over the architectureto compute the proxy, which can be time consuming on low-end hardware.
Additionally, theaforesaid pruning techniques, by definition, incorporate the final softmax projection layer in theirscore assessment.
Such an approach seems reasonable for CNNs dealing with a few classificationlabels, however, it can skew the evaluation for autoregressive Transformers dealing with a largeoutput vocabulary space.
To overcome these shortcomings, we introduce a zero-cost architectureranking strategy in the next section that outperforms the proposed low-cost proxies in terms of rankingprecision, is data free, and does not perform any forward/backward propagation. ▶
Decoder Parameter Count as a Proxy.
We empirically establish a strong correlation betweenthe parameter count inside decoder layers and final model performance in terms of validation perplexity.
We evaluate 200 architectures sampled uniformly at random from the search space of twoautoregressive Transformer backbones, namely, Transformer-XL and GPT-2.
These architectures arethen trained fully on WikiText-103 and One Billion Word (LM1B) datasets.
Full training requires 40 Kand 100 K steps on WikiText-103 and LM1B, respectively, and consumes over 25000 GPU-hours on
NVIDIA A100 and V100 nodes.
We compare the ranking obtained using our zero-cost proxy againstthe ground-truth ranking based on validation perplexity after full training in Figure 3.
On WikiText103, the zero-cost ranking using the decoder parameter count obtains a Spearman’s Rank Correlation(SRC) of 0 .
97 and 0 .
98 with full training for Transformer-XL and GPT-2 backbones, respectively.
SRC further increases to 0 .
99 for the more complex LM1B benchmark on both backbones.
Thisvalidates that the decoder parameter count is strongly correlated with final model performance,thereby providing a reliable zero-cost proxy for NAS.
Experiments
We conduct experiments to seek answers to the following critical questions:
1 How well can training-free proxies perform compared to training-based methods for estimatingthe performance of Transformer models? 2 How does model topology affect the performance of the proposed decoder parameter proxy?3 Which models are on the pareto-frontier of perplexity, latency, and memory for different hardware?4 Can our training-free decoder parameter count proxy be integrated inside a search algorithm toestimate the pareto-frontier?
How accurate is such an estimation of the pareto?We empirically answer questions 1 , 2 , and 3 in Sections 4.
4.3, and 4.4, respectively.
We further address question 4 in Appendix C where we show the pareto-frontier models extracted bythe decoder parameter count proxy are very close to the ground-truth pareto-frontier with an averageof 0 .
6% perplexity difference.
Additionally, we show the efficacy of the decoder parameter countproxy when performing search on different range of model sizes in Appendix C, Figure 10.
4.1 Experimental Setup
Please refer to Appendix B for information about the benchmarked datasets, along with the implementation details of our training, evaluation, and evolutionary search algorithm.
Backbones.
We apply our search on two widely used autoregressive Transformer backbones,namely, Transformer-XL [6] and GPT-2 [22] that are trained from scratch with varying architecturalhyperparameters.
The internal structure of these backbones are quite similar in that they containdecoder blocks with attention and feed-forward layers.
The difference between the backbones liesmainly in their dataflow structures; the models derived from Transformer-XL backbone adopt arecurrence methodology over past states, which coupled with relative positional encoding, allows formodeling longer term dependencies.
Performance Criteria.
To evaluate the ranking performance of various proxies, we first establisha ground-truth ranking of candidate architectures by training them until full convergence.
Thisground-truth ranking is then utilized to compute two performance criteria as follows:
▶ Common Ratio (CR): We define CR as the percentage overlap between the ground-truth ranking of architectures versus the ranking obtained from the proxy over the top- k % of architectures.
CR quantifies the ability of the proxy ranking to identify the top- k % architectures.
▶ Spearman’s Rank Correlation (SRC): We use this metric to measure the correlation between theproxy ranking and the ground-truth.
Ideally, the proxy ranking should have high correlation with theground-truth over the entire search space as well as high-performing candidate models.
4.2 How do training-free proxies perform compared to training-based methods?
In this section, we benchmark several proxy methods for estimating the rank of candidate architectures.
Specifically, we investigate three different ranking techniques, namely, partial training, low-costmethods, and number of decoder parameters.
▶ Partial Training.
We first analyze the relationship between validation perplexity after a shortenedtraining period versus that of full training for ranking candidate models.
We stop the trainingafter τ ∈ [1 . 25% , 87 .
5%] of the total training iterations needed for model convergence.
Figure 4demonstrates the SRC and CR of partial training with various τ s, evaluated on 100 randomlyselected models from the Transformer-XL backbone, trained on WikiText-103.
The horizontal axisdenotes the average time required for τ iterations of training across all sampled models.
Intuitively,a higher number of training iterations results in a more accurate estimate of the final perplexity.
Nevertheless, the increased wall-clock time prohibits training during search and also imposes theneed for GPUs.
Interestingly, very few training iterations, i.e., 1 . 25% , provides a very good proxyfor final performance with an SRC of > 0 .
9 on the entire population.
Our training-free proxy, i.e.,decoder parameter count, also shows competitive SRC compared to partial training.
▶ Low-cost Proxies.
We benchmark various low-cost methods introduced in Section 3.1 on 200randomly sampled architectures from the Transformer-XL backbone, trained on WikiText-103.
Figure 5 shows the SRC between low-cost proxies and the ground-truth ranking after full training. ground-truth ranking after full training of 200 randomly sampled Transformers. The decoder parameter count obtains the best SRC with zero cost.
We measure the cost of each proxy in terms offloating point operations (FLOPs).
As seen, theevaluated low-cost proxies have a strong correlation with the ground-truth ranking (even thelowest performing relu_log_det has > 0 .
SRC), validating the effectiveness of trainingfree NAS on autoregressive Transformers.
Thelower performance of relu_log_det can be attributed to the fundamental difference in muchhigher frequency of ReLU activations in CNNs,for which the method was originally developed,compared to Transformer-based architectures.
Our analysis on randomly selected models withhomogeneous structure also shows a strong correlation between the low-cost proxies and validationperplexity, with decoder parameter count outperforming other proxies (see Appendix D).
▶ Parameter Count.
Figure 6a demonstrates the final validation perplexity versus total numberof model parameters for 300 randomly sampled architectures from GPT-2 and Transformer-XLbackbones.
This figure contains two important observations: (1) the validation perplexity hasa downward trend as the number of parameters increase, (2)
The discontinuity is caused by thedominance of embedding parameters when moving to the small Transformer regime.
As seen,changing the adaptive input embedding [2] factor k ∈ { 1 , 2 , 4 } (shown with different colors in Figure 6a) varies the total parameter count without much influence on the validation perplexity.
The above observations motivate us to evaluate two proxies, i.e., total number of parameters andonly decoder parameter count.
Figures 6b,c demonstrate the CR and SRC metrics evaluated on the300 randomly sampled models divided into top- k % bins based on their validation perplexity.
As shown, the total number of parameters has a lower SRC with the validation perplexity, compared todecoder parameter count.
This is due to the masking effect of embedding parameters, particularly inthe Transformer-XL backbone.
The total number of decoder parameters, however, provides a highlyaccurate, zero-cost proxy with an SRC of 0 .
97 and 0 .
98 with the perplexity after full training on
Transformer-XL and GPT-2, respectively.
We further show the high correlation between decoderparameter count and validation perplexity for Transformer architectures with homogeneous decoderblocks in the supplementary material, Appendix D.
4.3 How does variation in model topology affect decoder parameters as a proxy?
The low-cost proxies introduced in Section 3.1, rely on forward and backward passes through thenetwork.
As such, they automatically capture the topology of the underlying architecture by relyingon the dataflow to compute the final proxy score.
The decoder parameter count proxy, however, istopology-agnostic.
In this section, we investigate the effect of topology on the performance of ourproposed decoder parameter count proxy.
Specifically, we seek to answer whether for a given decoderparameter count budget, the aspect ratio of the architecture, i.e., trading-off the width versus thedepth, can affect the final validation perplexity.
We define the aspect ratio of the architecture as d model (=width), divided by n layer (=depth).
This metricprovides a sense of how skewed the topology is and has been used in previous work which studyscaling laws for language models [14].
For a given decoder parameter count budget, we generateseveral random architectures from the GPT-2 backbone with a wide range of the width-to-depthaspect ratios 6 .
The generated models span wide, shallow topologies (e.g., d model = 1256 , n layer = 2 ) tonarrow, deep topologies (e.g., d model = 112 , n layer = 100 ).
Figure 7a shows the validation perplexityof said architectures after full training on WikiText-103 versus their aspect ratio.
The maximumdeviation (from median) of the validation perplexity is < 12 .
8% for a given decoder parametercount, across a wide range of aspect ratios ∼ 1 − 630 .
Our findings on the heterogeneous searchspace complement the empirical results by [14] where decoder parameter count largely determinesperplexity for homogeneous Transformer architectures, irrespective of shape (see Fig 5 in [14]).
We observe stable training when scaling models from the GPT-2 backbone up to 100 layers, withthe perplexity increasing only when the aspect ratio nears 1 .
Nevertheless, such deep models arenot part of our search space as they have a high latency and are unsuitable for lightweight inference.
For the purposes of hardware-aware and efficient Transformer NAS, decoder parameter count proxyholds a very high correlation with validation perplexity, regardless of the architecture topology asshown in Figure 7a.
We further validate the effect of topology on decoder parameter count proxy forthe Transformer-XL backbone in Figure 13 of Appendix
E. Our results demonstrate less than 7%deviation (from the median) in validation perplexity for aspect ratios ranging from ∼ 8 − 323 .
Note that while models with the same parameter count have very similar validation perplexities, thetopology in fact affects their hardware performance, i.e., latency (up to 2 . 8 × ) and peak memoryutilization (up to 5 . 5 × ), as shown in Figure 7b, 7c.
This motivates the need for incorporating hardwaremetrics in NAS to find the best topology.
[Figure]
4.4 Pareto-frontier models for various hardware platforms
The straightforward way to produce architectures of different latency/memory is varying the numberof layers (layer-scaling) [29, 6].
We therefore benchmark the Transformer-XL (base) and GPT-2(small) models with homogeneous layers ∈ [1 , 16] as our baseline as any practitioner is likely to do.
We compare architectures found by LTS with layer-scaled versions of the backbone models and showthat LTS models have better validation perplexity and/or lower latency and peak memory utilization(better pareto).
All baseline models are trained using the same setup 7 as the NAS-generated modelsfor fair comparison using the regimen in Table 2 of Appendix B.
We conduct LTS search on different target hardware.
LTS outputs a range of pareto-optimal architectures with various latency/memory/perplexity characteristics.
During search, we fix the adaptiveinput embedding factor to k = 4 to search models that are lightweight while ensuring nearly on-parvalidation perplexity with non-adaptive input embedding.
We compare the pareto-frontier architectures found by LTS with the baseline in Figure 8.
Here, all models are trained on the LM1B dataset(See Figure 12 in Appendix G for results on WikiText-103).
Note that the pareto-frontier search isperformed in a 3 -dimensional space, an example of which is enclosed in the Appendix, Figure 14.
For better visualization, in Figure 8 we plot 2 -dimensional slices of the pareto-frontier with validationperplexity on the y-axis, and a hardware performance metric (either latency or memory) on the x-axis.
As seen, in the low-latency regime, LTS consistently finds models that have significantly lowerperplexity compared to naive scaling of the baseline Transformer-XL or GPT-2.
On the TransformerXL backbone, LTS finds architectures with an average of 19 . 8% and 28 .
8% lower latency andmemory, while achieving similar perplexity compared to the baseline on ARM CPU.
Specifically,the perplexity of the 16 -layer Transformer-XL base can be replicated on the ARM device with alightweight model that is 1 .
6 × faster and utilizes 1 . 9 × less memory during execution.
On the Corei
CPU, the pareto-frontier models found by LTS are on average 25 . 8% faster and consume 30 . 0%less memory under the same validation perplexity constraint.
In this setting, LTS finds a model thatreplicates the perplexity of the 16 -layer Transformer-XL base while achieving 1 .
7 × faster runtimeand 1 .
9 × less peak memory utilization.
The savings are even higher on the GPU device, wherethe NAS-generated models achieve the same perplexity as the baseline with average 30 .
5% lowerlatency and 27 .
0% less memory.
Specifically, a LTS model with the same perplexity as the 16 -layer
Transformer-XL base has 2 .
5 × less latency and consumes 2 . {+
 ×+} less peak memory on TITAN Xp.
On the GPT-2 backbone, we observe that the NAS-generated models utilize on average 9 . 5% lessmemory while achieving the same validation perplexity and latency on an ARM CPU.
The benefits arelarger on Corei7 and TitanXP where the latency savings are pushed to 13 .
8% and 11 .
9% , respectively.
The peak memory utilization is also decreased by 9 . 7% and 12 .
9% , on average, compared to thebaseline scaled GPT-2s on Corei7 and TITAN Xp, respectively.
Notably, NAS finds new architectureswith the same perplexity as the 16 -layer GPT-2 with 1 . 3 × , 1 . 4 × faster runtime and 1 . × less peakmemory utilization on Corei7 and TITAN Xp.
Our heterogeneous search-space allows us to find a better parameter distribution among decoderlayers.
Therefore, LTSelivers architectures with better performance in terms of perplexity, while Architectural parameters for models shown here are detailed in Appendix H. reducing both the latency and memory utilization when compared to the homogeneous baselines.
We provide the architecture of all baseline and LTS models shown in Figure 8 in Tables 3-6 of Appendix H.
Search Efficiency.
The main component in LTS search time is the latency/peak memory utilizationmeasurement for candidate architectures since evaluating the model perplexity is instant using thedecoder parameter count.
Our search only takes 0 .
9 , 2 . 6 , and 17 .
2 hours on a TITAN
Xp GPU, Corei
CPU, and an ARM core, respectively.
To provide more context into the timing analysis, full trainingof even one 16 -layer Transformer-XL base model on LM1B using a machine with 8 × NVIDIA[-V
-]
V100 GPUs takes 15 . 8 hours.
Using decoder parameters as a proxy significantly reduces our searchruntime, allowing it to finish in a few hours on commodity hardware.
Once the pareto-frontier modelsare identified, the user can pick a model based on their desired hardware constraints and fully train iton the target dataset.
LTS is an alternate paradigm to that of training memory-hungry supernets; oursearch can run directly on target device (not just memory/latency measurement) with GPUs neededonly for training the chosen pareto-frontier model after search.
In Table 1 we study the ranking performance of partial training ( 500 ) steps against the decoderparameters proxy for evaluating 1200 architectures from the Transformer-XL backbone during LTSsearch.
Astonishingly the decoder parameters proxy gets higher SRC against partial training.
Thisproxy can hence be effectively used in the inner loop of search for NAS.
[Figure]
[Table]
Limitations and Future Work
Decoder parameter count provides a simple yet accurate proxy for ranking autoregressive models.
This should serve as a strong baseline for future works on Transformer NAS.
Our focus is entirelyon autoregressive, decoder-only transformers.
We therefore, study perplexity as the commonly usedmetric for language modeling tasks.
Nevertheless, recent literature on scaling laws for Transformerssuggest a similar correlation between parameter count and task metrics may exist for encoder only(BERT-style)
Transformers or encoder-decoder models used in neural machine translation (NMT) [11].
Additionally, recent findings [26] show specific scaling laws exist between model size and downstreamtask metrics, e.g., GLUE [30].
Investigating such scaling laws for NAS on heterogeneous BERT-styleor NMT models with new metrics is an important future avenue of research.
Checklist
The checklist follows the references.
Please read the checklist guidelines carefully for information onhow to answer these questions.
For each question, change the default [TODO] to [Yes] , [No] or[N/A] .
You are strongly encouraged to include a justification to your answer , either by referencingthe appropriate section of your paper or providing a brief inline description.
For example:
• Did you include the license to the code and datasets?
[Yes] See Section XXX.
• Did you include the license to the code and datasets?
[No] The code and the data are
proprietary.• Did you include the license to the code and datasets?
[N/A]
Please do not modify the questions and only use the provided macros for your answers.
Note that the
Checklist section does not count towards the page limit.
In your paper, please delete this instructionsblock and only keep the Checklist section heading above along with the questions/answers below.
For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper’scontributions and scope?
[Yes]
(b) Did you describe the limitations of your work?
[Yes] See Sections 4.3, 5, Appendix E(c)
Did you discuss any potential negative societal impacts of your work?
[Yes] SeeAppendix I
(d) Have you read the ethics review guidelines and ensured that your paper conforms to
them?
[Yes]
If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results?
[N/A]
(b) Did you include complete proofs of all theoretical results?
[N/A]
If you ran experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main experimental results (either in the supplemental material or as a URL)?
[No] All trainedmodels will be open-sourced upon acceptance.
(b) Did you specify all the training details (e g., data splits, hyperparameters, how theywere chosen)?
[Yes] See Appendix B(c)
Did you report error bars (e g., with respect to the random seed after running experi-
ments multiple times)?
[N/A](d)
Did you include the total amount of compute and the type of resources used (e.g., typeof GPUs, internal cluster, or cloud provider)?
[Yes]
If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...
[Table]
(a) If your work uses existing assets, did you cite the creators?
[Yes] See Section B
[Table]
(b) Did you mention the license of the assets?
[N/A] All used code are open-source and
[Table]
publicly available
[Table] (c) Did you include any new assets either in the supplemental material or as a URL?
[No]
[Table]
(d) Did you discuss whether and how consent was obtained from people whose data you’re
[Table]
using/curating?
[N/A]
[Table] (e) Did you discuss whether the data you are using/curating contains personally identifiable [Table]
information or offensive content?
[N/A]
[Table]
If you used crowdsourcing or conducted research with human subjects...
[Table]
(a) Did you include the full text of instructions given to participants and screenshots, if
[Table]
applicable?
[N/A]
[Table]
(b) Did you describe any potential participant risks, with links to Institutional Review
[Table]
Board (IRB) approvals, if applicable?
[N/A]
[Table]
(c) Did you include the estimated hourly wage paid to participants and the total amount [Table] spent on participant compensation?
[N/A] [Table]