ETAS : Zero-Shot Transformer Architecture Search via Network Trainability and Expressivity (2024)

Abstract: Transformer Architecture Search (TAS) methods aim to automate searching for the optimal Transformer architecture configurations for a given task. However, they are impeded by the prohibitive cost of evaluating Transformer architectures. Recently, several Zero-Shot TAS methods have been proposed to mitigate this problem by utilizing zero-cost proxies to evaluate Transformer architectures without training. Unfortunately, they are limited to specific computer vision or natural language processing tasks. Nonetheless, most of them are developed based on empirical observations and lack theoretical guarantees. To solve this problem, we develop a new zero-cost proxy called NTSR that combines two theoretically-inspired indicators to measure the trainability and expressivity of Transformer networks separately. We then integrate it into an effective regularized evolution framework called ETAS to demonstrate its efficacy on various tasks. The results show that our proposed NTSR proxy can consistently achieve a higher correlation with the true performance of Transformer networks on both computer vision and natural language processing tasks. Further, it can significantly accelerate the search process for finding the best-performing Transformer architecture configurations.

## introduction transformer networks @xcite @xcite have attracted tremendous interest over recent years due to their effectiveness in learning long-range dependencies in data and superior performance across various tasks. they have gradually replaced traditional neural networks, such as convolutional neural networks (cnns) and recurrent neural networks (rnns), in a variety of domains including natural language processing (nlp) @xcite , computer vision (cv) @xcite , * corresponding author speech signal processing @xcite , and healthcare @xcite . recently, the transformer architecture has become the de facto backbone for most large language models (llms) @xcite . in real applications, it is often necessary to adjust the transformer architecture configurations according to the specific tasks @xcite , such as the depth of the network, the number of attention heads, embedding dimension, and the inner dimension of the feed-forward layer. however, manually tuning these parameters requires repeated refinement with expert experience, which is timeconsuming and computationally expensive. to solve this problem, various transformer architecture search (tas) methods @xcite @xcite have been proposed, which automate the search for the optimal transformer architecture configurations for a given task and data. the current popular tas methods include reinforcement learning @xcite , evolutionary search @xcite , one-shot @xcite , and predictorbased search @xcite . however, during the search process, they still demand a high computational cost to evaluate several hundred or thousands of architectures. training a transformer network can take hours or even days, thus hindering the practical application of tas @xcite . recently, zero-shot neural architecture search (nas) methods @xcite have attracted much attention as they design zero-cost proxies to estimate the performance of a network at the initialization stage. they can quickly evaluate the performance of a network in a few seconds by computing statistics from a single forward/backward propagation pass of the network with a minibatch of data at initialization @xcite . nevertheless, @xcite have shown the majority of existing zero-cost prox-ies are specifically designed for the cnn search spaces (e.g.,nas-bench-101 @xcite , nas-bench-201 @xcite , nas-bench-301 @xcite , transnas-bench-101 @xcite ) and perform worse on the transformer search space. they leverage the characteristics of transformer networks and design a dss zero-cost proxy that estimates synaptic diversity of multi-head self-attention (msa) and synaptic saliency of the multilayer perceptron (mlp) in the transformer network to rank its performance in a vision transformer (vit) search space. later, @xcite introduce a zero-cost proxy that measures the complexity of manifold propagation through vit to estimate how complex a function can be approximated by a vision transformer network. @xcite choose the number of decoder parameters in auto-regressive transformers as a zero-cost proxy for the perplexity of the language model without the need for any model training. unfortunately, current zero-cost proxies for transformer architectures are specifically designed for certain computer vision or natural language processing tasks. for example, existing zero-cost proxies for transformers like dss @xcite and tvt @xcite are specially designed for vision transformer (vit) encoder-only architectures. they show great performance on multiple vision tasks, but they perform relatively worse on the most used gpt auto-regressive architecture in nlp applications @xcite . on the other hand, lts @xcite show that decoder parameters serve as a good proxy for the performance of gpt-2 architectures on multiple nlp tasks. however, they perform relatively worse in vit architecture @xcite as there exists a large amount of parameter redundancy in these structures, leading to overparameterized networks that are prone to overfitting on small-scale datasets. nonetheless, most of them are designed based on empirical observations and lack theoretical assurances. can we design a theoretically-inspired zero-cost proxy applicable to multiple vision and language tasks? to this end, we propose a novel zero-cost proxy called ntsr that combines two theoretically-inspired indicators to measure the trainability and expressivity of transformer networks. in particular, based on the theoretical underpinnings of deep neural network training, we design the ntkt metric that utilizes the trace of the mean neural tangent kernel (ntk) to quantify the trainability of transformer networks. meanwhile, we design another sept metric that utilizes the upper bound of separation rank induced by the transformer network to measure the capacity of transformer networks to represent input dependencies. to demonstrate the effectiveness of our proposed ntsr zero-cost proxy for transformer networks, we compare it to other popular zero-cost proxies in multiple search spaces. the results show that ntsr can consistently achieve a higher correlation with the true performance of transformer networks on both computer vision and natural language processing tasks. to investigate the ability of our proposed ntsr zero-cost proxy to accelerate searching for the best-performing transformer network, we further integrate it into an effective regularized evolution framework called etas. the results show that ntsr can significantly speed up finding the best-performing transformer network. zero-cost proxies. in recent years, various types of zero-cost proxies @xcite @xcite have been proposed to rank the performance of networks at the initialization stage. one type consists of parametric saliency-based zero-cost proxies @xcite . they score the whole network by summing the changes in the saliency metric when a specific parameter of the network is removed. @xcite adopt several pruning-atinitialization metrics, including snip, grasp, synflow, and fisher, as parameter-level saliency-based zero-cost proxies for the network performance with a minibatch of data at initialization. another type is network expressivity-based zero-cost proxies @xcite . they estimate the network performance by the expressivity of the network. @xcite develop a heuristic naswot metric that utilizes the correlation of network activations between different data at initialization. meanwhile, @xcite select the number of linear activated regions represented by the network to measure the expressivity of a network. additionally, there exists a type of zero-cost proxies inspired by deep learning theory @xcite . @xcite choose the condition number of neural tangent kernel (ntk) to measure the trainability of networks and show it is negatively correlated with network performance. later, @xcite develop a label-agnostic and data-agnostic nasi metric that leverages the trace of ntk as an indicator for network performance. more recently, @xcite found that the mean value and standard deviation of gradients across different samples affect the training convergence of networks. based on the generalization theory of deep neural networks, they propose a theoretically-inspired zero-cost proxy called zico that considers both the mean value and standard deviation of gradients in each layer of the network. meanwhile, @xcite use the minimum eigenvalue of the pearson correlation matrix upon each layer of the feature maps to indicate the convergence of the network. however, most of the existing zero-cost proxies are specially designed for cnn networks, whose internal architecture is distinctly different from the transformer network. @xcite have shown that the current zero-cost proxies for cnn cannot generalize well to the transformer search space. to adapt well for the transformer network, we develop a novel zero-cost proxy called ntsr that combines two theoretically-inspired indicators for measuring the trainability and expressivity of the transformer network. ## method in the context of tas, the goal of zero-cost proxy is to accurately estimate the ranking of a transformer network's performance at initialization. in fact, a good neural network architecture should have good trainability (i.e., how fast a network can converge via a gradient descent algorithm) and high expressivity (i.e., how complex functions a network can represent) @xcite . in sections 3.1 and 3.2, we design two theoretically-inspired indicators to reflect the trainability and expressivity of transformer networks, separately. we then propose the ntsr zero-cost proxy that combines these two important indicators to measure the trainability and expressivity of a given transformer network, and integrate it into an evolutionary search framework called etas to find the best-performing transformer network architecture in section 3.3. ## experiments and discussion in this section, we choose to search for the most popular vision transformer (vit) and gpt-2 architectures in the computer vision and natural language processing domains, respectively. we compare our proposed ntsr zero-cost proxy with several zero-cost proxies for transformers, including dss @xcite , length distortion (ld) @xcite , and conventional zerocost proxies for cnns like snip @xcite , grad_norm @xcite , naswot @xcite , zico @xcite , and meco @xcite . @xcite have demonstrated that the number of flops and parameters serve as competitive proxies for network performance compared to most zero-cost proxies under the cnn search space, we also add them for comparison. to perform an ablation study on our proposed ntsr proxy, we include ntkt and sept alone as zero-cost proxies for comparison. here, we choose the widely used spearman's ρ (krishnakumar et al., 2022b) and kendall's τ @xcite rank correlation metrics to evaluate their predictive ability. ## limitation and conclusion in this study, we propose a novel zero-cost proxy for transformer networks called ntsr to evaluate the performance of transformers at initialization. compared to other popular zero-cost proxies for transformer networks, our proposed proxy is designed based on deep neural network learning theory in terms of the trainability and expressivity of transformer networks. it achieves better performance on both nlp and cv tasks. however, our proposed method is limited to decoder-only transformer architectures. in the future, we hope our proposed proxy can extend to other complex types of transformer networks like bert @xcite and t5 @xcite .