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.

heading: 1 Introduction
text: Transformer networks (Li et al., 2022a; Zhou et al., 2022; Chitty-Venkata et al., 2022) 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) (Javaheripi et al., 2022), Computer Vision (CV) (Chen et al., 2022),
* Corresponding Author
speech signal processing (Chitty-Venkata et al., 2022), and healthcare (Chitty-Venkata et al., 2022). Recently, the Transformer architecture has become the de facto backbone for most large language models (LLMs) (Zhao et al., 2023; Tornede et al., 2024). In real applications, it is often necessary to adjust the Transformer architecture configurations according to the specific tasks (Javaheripi et al., 2022), 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 (Lin et al., 2021; Chitty-Venkata et al., 2022; Wei et al., 2024; Chen et al., 2023) 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 (Zhu et al., 2021), evolutionary search (Real et al., 2019; Zhou et al., 2024), one-shot (Li et al., 2021; Chen et al., 2021b), and predictorbased search (Wei et al., 2023; Ning et al., 2023). 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 (Wei et al., 2024). Recently, zero-shot Neural Architecture Search (NAS) methods (Li et al., 2024) 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 (Abdelfattah et al., 2021). Nevertheless, Zhou et al. (2022) have shown the majority of existing zero-cost prox-
6780
ies are specifically designed for the CNN search spaces (e.g.,NAS-Bench-101 (Ying et al., 2019), NAS-Bench-201 (Dong and Yang, 2020), NASBench-301 (Zela et al., 2022), TransNAS-Bench101 (Duan et al., 2021)) 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, Chen et al. (2022) 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. Javaheripi et al. (2022) 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 (Zhou et al., 2022) and TVT (Wei et al., 2023) 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 (Chen et al., 2022). On the other hand, LTS (Chen et al., 2022) 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 (Zhou et al., 2024; Zhou and Zhu, 2024) 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.

heading: 2 Related Work
text: Transformer Architecture Search. The goal of Transformer Architecture Search (TAS) is to search for the best-performing Transformer configurations in an automated way for a given task. Recently, various TAS methods (Li et al., 2022b; ChittyVenkata et al., 2022; Ye et al., 2024) have been proposed. One of the most popular methods is evolution-based Transformer architecture search (So et al., 2019, 2021; Su et al., 2022). They use an evolutionary search algorithm to find the optimal Transformer architecture within a given target budget. Another classic method is reinforcement learning (Zhu et al., 2021; Zhu, 2021), which uses a controller to sample high-quality Transformer architectures and updates the controller using the network’s performance as a reward. However, current evolutionary search and reinforcement learning approaches need to collect a substantial number of network samples to train, which is costly and time-consuming. To reduce the search cost of TAS, one-shot methods (Chen et al., 2021a,b; Li et al., 2022a; Wang et al., 2023) have been proposed. They only train a huge supernetwork and obtain the performance of the sampled subnetwork through weight sharing. Nevertheless, training a huge supernetwork is non-trivial, and the memory consumption of a Transformer supernetwork increases with hidden size and runs out even with small values (Chitty-Venkata et al., 2022). Predictor-based methods (White et al., 2021; Shen et al., 2023) train
a surrogate model using a certain number of Transformer architecture-accuracy pairs and then use it to estimate the performance of unseen networks. Overall, the current popular TAS methods still need fully training massive architectures to find the optimal network, which demands a high computational cost. This highlights the great potential of developing zero-shot methods to replace the expensive training process in NAS with zero-cost proxies.
Zero-Cost Proxies. In recent years, various types of zero-cost proxies (Krishnakumar et al., 2022a; Li et al., 2024; Lee and Ham, 2024) have been proposed to rank the performance of networks at the initialization stage. One type consists of parametric saliency-based zero-cost proxies (White et al., 2022). They score the whole network by summing the changes in the saliency metric when a specific parameter of the network is removed. Abdelfattah et al. (2021) 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 (Tu et al., 2022). They estimate the network performance by the expressivity of the network. Mellor et al. (2021) develop a heuristic NASWOT metric that utilizes the correlation of network activations between different data at initialization. Meanwhile, Chen et al. (2021c) 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 (Zhou and Zhu, 2024; Shu et al., 2022b). Chen et al. (2021c) 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, Shu et al. (2022a) develop a label-agnostic and data-agnostic NASI metric that leverages the trace of NTK as an indicator for network performance. More recently, Li et al. (2023) 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, Jiang et al. (2023) 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. Zhou and Zhu (2024) 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.

heading: 3 Method
text: 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) (Chitty-Venkata et al., 2022; Chen et al., 2021c). 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.

heading: 3.1 Trainability of Transformer Network
text: With the rapid development of deep learning theory on neural networks, Neural Tangent Kernel (NTK) has emerged as an effective tool for characterizing the training dynamics of infinitely wide (Jacot et al., 2021) or finite wide (Novak et al., 2022) deep networks. It solved a classic question of “how does training of neural networks work so well despite being highly nonconvex?" (Yang, 2020). Lee et al. (2019) have demonstrated that under the large width limit and constant NTK assumption, the predictions of a neural network evolving like a linear model throughout gradient descent training.
NTK of Transformer. Assume there exists a batch of M sequences with T tokens X = {xα,1, · · · ,xα,T }Mα=1 and a standard depth-L Transformer network with one input embed-
ding layer and L transformer blocks. The output of Transformer network at the L-th block is {zLα,1, · · · , zLα,T }Mα=1. After that, we define the NTK of a Transformer network K(X ,X ) ∈ RMdLz ×T×MdLz ×T , each element of the 4-dimensional NTK tensor is:
Kα1,t1,α2,t2 = ∇wzLα1,t1 ( ∇wzLα2,t2 )⊺ , (1)
where α1, α2 ∈ X are pairs of inputs sampled from training batch X . t1, t2 ∈ [T ] is the token index. Theorem 3.1 (Yang (2020)). Let f be a neural network of standard architecture with scalar output and randomly initialized weights ω ∼ N ( 0, σ2ω ) . If it satisfies Condition 1 in Yang (2020) and its nonlinearities have polynomially bounded weak derivatives, then its NTK Θ converges almost surely, over any finite set of inputs, to a deterministic kernel Θ0, i.e., Θ a.s.−−→ Θ0 as its widths go to infinity.
This theorem indicates the NTK of a standard Transformer network at initialization K(X ,X ) has a well-defined infinite-width limit. We can use it to predict the training convergence of the network. As shown in Equation (14), we can estimate the training convergence of the network through NTK at initialization. If Θ0 can be represented through its eigenvectors and corresponding eigenvalues λi, then it can be inferred that the eigenvectors of Θ0 coincide with those of e−ηΘ0t with a transformation of eigenvalues e−ηλit. This observation reveals that the network’s convergence rate is intimately connected to eigenstructures of Θ0. Consequently, a network with a NTK characterized by a greater total sum of eigenvalues, i.e., ∑ i λi, is likely to achieve faster convergence and a lower loss. However, directly computing ∑ i λi of the 4- dimensional NTK tensor K(X ,X ) of a Transformer network is extremely challenging. This difficulty arises due to the dynamic nature of the network connections in Transformers, where the output at each token index focuses on different parts of the input sequence, as shown in Equation (7), resulting in distinct outputs for each token index. To solve this problem, we take the mean of the NTK tensor K(X ,X ) over the token indexes, which is defined as K̄(X ,X ) ∈ RMdLz ×MdLz , where each element is:
K̄α1,α2 = 1
T 2
T∑
t1=1
T∑
t2=1
∇wzLα1,t1 ( ∇wzLα2,t2 )⊺ .
(2)
Through the above operation, the four-dimensional NTK tensor K(X ,X ) is reduced to a standard twodimensional matrix K̄α1,α2 . Besides, the matrix is symmetric and positive semi-definite, and the sum of eigenvalues is equivalent to the trace of K̄(X ,X ). Calculating the trace of a matrix offers computational efficiency, effectively bypassing the complexities involved in eigenvalue determination. Consequently, we use the trace of K̄α1,α2 as an indicator called NTKT for the training convergence of the network, which is defined as:
NTKT(f) = ∥∥K̄(X ,X ) ∥∥ tr = M∑
i=1
∥∥K̄αi,αi ∥∥
tr
= 1
T 2
M∑
αi=1
T∑
ti=1
dLz∑
j=1
∥∥∇wzLαi,ti,j ∥∥2 2 .
(3)
Theorem 3.2. Assume a standard Transformer net f with randomly initialized weights ω ∼ N ( 0, σ2ω ) . Under vanilla stochastic gradient descent (SGD) optimizer and the first-order Taylor expansion, the mean of outputs over tokens at time step s + 1 satisfies: f̄s+1(X ) − f̄s(X ) = −ηsK̄(X ,X )L′ ( f̄s (X ) ) , where ηs is the learning rate at step s. Then, NTKT(f) is positively correlated with the convergence rate of network f .
The proof of this theorem can be found in Appendix A.2. This theorem shows that NTKT provides a reasonable estimate for the training convergence of the network. In general, a larger NTKT score indicates that the corresponding network will converge faster and learn more efficiently.

heading: 3.2 Expressivity of Transformer Network
text: The success of Transformer networks largely relies on attention mechanism to learn complex dependencies among input sequences for different tasks. Intuitively, the stronger a Transformer network can model dependencies between inputs, the greater its capability to represent input information. Is there a quantitative metric to measure the expressivity of a Transformer network at initialization? To solve this problem, we introduce separation rank as a metric to quantify its ability for modeling dependencies between inputs of a Transformer network.
Separation Rank for Transformer. As shown in Equation (15), the separation rank directly reflects the core attention mechanism in Transformer networks. The self-attention layer dynamically learns
to correlate any inter-dependent subsets of the inputs (Levine et al., 2020). When a transformer network learns more dependencies between inputs, the separation rank induced by this network will have a higher value. Levine et al. (2020) have demonstrated a tight bound on the separation rank of Transformer networks with respect to the dependence on depth and width. Further, they leverage this bound to determine the optimal depth-to-width ratio for a given Transformer network size.
Theorem 3.3 (Levine et al. (2020)). Let f i,t,L,dx,H be the i-th scalar output at the t token index of a standard depth-L Transformer net f with dimension dx and the number of heads H per layer. Then, its separation rank w.r.t. balanced partitions, which obey A ∪ B = [T ], |A| = |B| = T/2, is invariant to identity of the partition, i.e., A ∪ B = [T ], Ã ∪ B̃ = [T ], such that |A|, |B|, |Ã|, |B̃| = T/2: sep(A,B)(f i,t,L,dx,H) = sep(Ã,B̃)(f i,t,L,dx,H).
This theorem reveals that the separation rank induced by the standard Transformer network remains consistent under different balanced partitions of inputs. Next, we will omit the specification of any balanced partition of inputs, denoting it as sep(f i,t,L,dx,H). Wies et al. (2021) further demonstrate the existence of an embedding rank bottleneck in the expressivity of the Transformer network. They showed that log(sep(f i,t,L,dx,H)) = Õ (L ·min {r, dx}), where r is the embedding rank defined as r = rank(WV ).
However, the current study mainly focuses on the separation rank of the Transformer network with the same dimension dx per layer. In this study, we extend the separation rank to the Transformer network where each block l has different dlz and dlin dimensions. This makes it more difficult to analyze. To solve this problem, we follow the relaxations by Wies et al. (2021) and Levine et al. (2020). We put all the position-wise feed-forward layers at the end since the feed-forward operation does not mix different locations and learn the dependencies between inputs. We then remove all the normalization layers and omit the ReLU and softmax non-linearities. The reasons for doing this can be referred to in Wies et al. (2021) and Levine et al. (2020). Though these relaxations are shown to weaken the overall network performance, they are much less pertinent to the ability of the core self-attention mechanism to model dependencies between different places at the input. Consequently,
the multi-head attention operation of each block l in Equation (8) can be simplified as:
f l+1t = T∑
t′=1
H∑
h=1
〈 W q,l,hf lt ,W k,l,hf lt′ 〉 W o,l,hW v,l,hf lt′ .
(4)
By forward propagating the above operations layerby-layer, we can obtain the upper bound on the separation rank of the Transformer network with different dimensions dlz per layer. Theorem 3.4. Let fLt be the output at the t token index of a standard depth-L Transformer net f with dimension dlz and the number of heads H per layer. Then, its separation rank w.r.t. balanced partitions satisfies log(sep(fLt )) ≤ log( ∑L l=1 d l z) + log( ∑L
l=1 1−(dlz)T+1
1−dlz ).
The proof of this theorem is shown in Appendix A.2. This theorem shows that the separation rank of the Transformer net f can be upper bounded by a constant, which is related to the depth of the network L, the number of tokens T , and the sum of dimensions ∑L l=1 d l z . Thus, we use this bound to measure the amount of dependencies modeled by the Transformer net. We name this measure as SEPT, which is defined as:
SEPT(f) = log(
L∑
l=1
dlz) + log(
L∑
l=1
1− (dlz)T+1 1− dlz )
(5)
SEPT measures the Transformer network’s ability to model dependencies between different places of the input. In general, a larger SEPT score indicates that the corresponding Transformer network has a stronger expressivity regarding different dependencies between inputs.

heading: 3.3 Zero-Shot Transformer Architecture Search via Network Trainability and Expressivity
text: NTSR. How can we combine the two theoretically-inspired indicators, NTKT and SEPT, to find the best-performing Transformer network at initialization? To answer this question, we propose a new zero-cost proxy called NTSR to make a trade-off between the network’s trainability and expressivity. Note that the magnitudes of NTKT and SEPT scores may differ significantly, and their ranges are unknown before computing on the whole search space. Hence, we cannot
directly normalize NTKT and SEPT scores before searching the network. To avoid this problem, instead of using the numerical values of NTKT and SEPT, we combine the relative ranking of NTKT and SEPT by comparing the sampled set of architectures, which is defined as:
NTSR(f) = 1
2 (R(NTKT(f)) + R(SEPT(f))),
(6)
where R(NTKT(f)) and R(SEPT(f)) represent the relative ranking of the network f converted from the NTKT(f) and SEPT(f) among the sampled set of architectures. In general, a higher rank of NTSR indicates that the network has better trainability and expressivity.
ETAS. As described above, our proposed NTSR zero-cost proxy provides a reasonable estimation of the network’s performance. The next major question is how to construct an efficient NAS framework on top of it? To solve this problem, we integrate it into a popular regularized evolution framework called ETAS. We first randomly sample N0 network architectures from the search space and then pick the top n0 networks according the ranking of NTSR as the initial parent population to warmup the whore algorithm. Through warmup, we can find a better local initial population, thus potentially expediting the discovery of the optimal network. Note that the sample size N0 is much larger than the number of networks we can afford to train. After that, we generate Nm candidate architectures by mutating the parent architectures at each iteration m and then select top nm architectures from the current Nm candidate architectures according the ranking of NTSR. We then evaluate the selected architectures and update the parent
population. The algorithm details of our proposed ETAS framework is given in the Appendix A.3.

heading: 4 Experiments and Discussion
text: 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 (Zhou et al., 2022), Length Distortion (LD) (Chen et al., 2022), and conventional zerocost proxies for CNNs like snip (Abdelfattah et al., 2020), grad_norm (White et al., 2022), NASWOT (Mellor et al., 2021), zico (Li et al., 2023), and MeCo (Jiang et al., 2023). Since Ning et al. (2021) 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 τ (Ning et al., 2021) rank correlation metrics to evaluate their predictive ability.

heading: 4.1 Searching for ViT
text: To make a fair comparison, we employ the same search space of AutoFormer (Chen et al., 2021a), which searches the key components of the ViT architecture including embedding dimension, Q-K-V dimension, number of heads, MLP ratio, and network depth. It contains more than 1.7× 1016 candidate architectures covering three common ranges of model size, i.e., tiny (4-9 M), small (14-34 M),
and base (42-75 M). We randomly sample 1000 networks for each network setup since it is computationally infeasible to evaluate all networks of the large AutoFormer search space. After that, we compute the rank correlation between these zero-cost proxy scores and the true accuracy of these sampled networks on the ImageNet-1K dataset. Table 1 shows the performance of various zero-cost proxies across three setups, i.e., tiny, small, and base. The results show that our proposed NTSR proxy can achieve similar ranking correlation with the oneshot AutoFormer method and perform better than other zero-cost proxies.
To demonstrate the effect of the NTSR proxy on accelerating TAS, we integrate it into the ETAS framework. Here, we choose to incorporate the top5 zero-cost proxies from Table 1 into our proposed ETAS framework to evaluate their performance on speeding up the regularized evolution algorithm in discovering the best ViT architecture. The implementation details and experimental settings of these methods are summarized in Appendix A.4.1. Figures 1(a) to 1(c) show the best-found valid loss over number of iterations of various methods. The test loss of various methods on the test set of ImageNet1k dataset is shown in Table 2. The results reveal that our proposed NTSR zero-cost proxy can achieve higher rank correlation across three setups. It can find a Transformer network architecture with lower validation loss using fewer iterations.
To further demonstrate the flexibility of our proposed NTSR proxy, we compare it with the recently proposed TVT (Wei et al., 2023) zero-cost proxy for ViT on three tiny datasets (CIFAR-100, Flowers, Chaoyang) of the AutoFormer search space.
Here, we randomly sample 100 ViT architectures from the AutoFormer search space and compute the Kendall ranking correlation and Spearman ranking correlation between distillation accuracy and different zero-cost proxies. The results reveal that our proposed NTSR zero-cost proxy still performs better on these three tiny datasets.

heading: 4.2 Searching for GPT-2
text: The GPT-2 search space used in this section is largely based on the design of the LTS (Javaheripi et al., 2022) search space on the WikiText-103 (Merity et al., 2017) and One Billion Word (LM1B) (Chelba et al., 2014) datasets. It consists of the number of layers (n_layer ∈ {2, . . . , 16}|1), number of attention heads (n_head ∈ {2, 4, 8}), decoder output dimension for each layer (d_model ∈ {128, . . . , 1024}|64), inner dimension of the feed forward network for each layer (d_inner ∈ {256, . . . , 4096}|64), and embedding dimension (d_embed ∈ {128, 256, 512}). Unlike the LTS search space, we fix the adaptive input embedding factor to k = 4 to approximate the standard Transformer network with non-adaptive input embedding. We add a constraint that d_inner must be larger than 2d_model to avoid training collapse of the network. To quantify the distance between the proxy and ground truth Pareto front, we use the mean average perplexity difference (MAPD) metric used in LTS, which is defined as: davg = 1 N ∑N i=1 |pi−pgt,i| pgt,i
, where pi denotes the ith point on the proxy Pareto front and pgt,i is the closest ground truth Pareto front to pi.
For each dataset, we randomly sample 500 networks and compute zero-cost proxies for each network. We train each GPT-2 model from scratch to obtain its true validation perplexity. After that, we compute rank correlation between these proxy scores and true validation perplexity of these sampled networks on the WikiText-103 and One Billion Word datasets. Table 4 shows the performance of various zero-cost proxies. To demonstrate the effect
of the NTSR proxy on speeding up TAS, we integrate it into the ETAS framework. Here, we choose to incorporate the top-5 zero-cost proxies from Table 4 into our proposed ETAS framework to evaluate their impact on speeding up the regularized evolution algorithm in discovering the best GPT-2 architecture on the WikiText-103 dataset. The implementation details and experimental settings of these methods are summarized in Section 4.2.
The validation perplexity of various methods on GPT-2 search space is shown in Table 5. The results reveal that our proposed NTSR proxy has a higher rank correlation with the true validation perplexity. It can significantly speed up finding betterperforming Transformer network architectures with fewer iterations. In contrast, the zero-cost proxies for Transformers like DSS and LD perform worse on this GPT-2 search space. This may be because they are specially designed for the ViT transformer architecture and thus generalize worse on the GPT Transformer network. Furthermore, our proposed NTSR proxy is closer to the true Pareto front compared to other zero-cost proxies.

heading: 4.3 Ablation Studies
text: As shown above, for some tasks like the WikiText103 dataset within the GPT-2 search space (see Table 4), NTKT performs better than SEPT and achieves lower validation perplexity. However, for certain tasks like the ImageNet-1k dataset within the AutoFormer small search space (see Table 1), SEPT performs better than NTKT. Our proposed NTSR zero-cost proxy combines the two theoretically-inspired indicators, NTKT and SEPT, and obtains consistently better performance across multiple tasks. To further analyze the combination effect of these two indicators, we perform ablation experiments on the combination coefficients of NTKT and SEPT (Equation (6)) on the WikiText103 dataset within GPT-2 search space and the ImageNet-1k dataset within the AutoFormer small search space. The results are shown in Section 4.3. The results demonstrate the advantage of combining NTKT and SEPT indicators equally. This may be because combining NTKT and SEPT equally allows them to complement each other.

heading: 5 Limitation and Conclusion
text: In this study, we propose a novel zero-cost proxy for Transformer networks called NTSR to evaluate the performance of Transformers at initializa-
tion. 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 (Devlin et al., 2019) and T5 (Raffel et al., 2020).

heading: Acknowledgements
text: We thank the anonymous reviewers for their valuable and constructive suggestions and comments. This work is supported by the Beijing Natural Science Foundation (No.4222029); the National Natural Science Foundation of China (N0.62076234); the National Key Research and Development Project (No.2022YFB2703102); the “Intel ligent Social Governance Interdisciplinary Platform, Ma jor Innovation & Planning Interdisciplinary Platform for the “Double-First Class” Initiative, Renmin University of China”; the Beijing Outstanding Young Scientist Program (NO.BJJWZYJH012019100020098); the Public Computing Cloud, Renmin University of China; the Fundamental Research Funds for the Central Universities, and the Research Funds of Renmin University of China (NO.2021030199), the Huawei-Renmin University joint program on Informa tion Retrieval: the Unicom Innovation Ecological Coopera tion Plan; the CCF-Huawei Populus Grove Fund.

heading: A Appendix

heading: A.1 Background

heading: A.1.1 Standard Transformer
text: Assume a standard depth-L Transformer network has one input embedding layer and L transformer blocks. The input of the Transformer network is a sequence of T tokens {xt ∈ [V ]}Tt=1, where V is the number of vocabulary tokens. The embedding
layer transforms the input sequence {xt ∈ [V ]}Tt=1 to T sequenced d1x-dimensional vectors z 1 t , t ∈ [T ], which is defined as z1t = WV xt + pt, where WV ∈ Rd 1 x×V is the learned word-embedding matrix, and pt is the positional embedding vector. After that, z1t is recursively transformed into T sequenced dlx-dimensional vectors z l t, t ∈ [T ], l ∈ [L] := {1, . . . , L} through L transformer blocks. Each transformer block consists of two sublayers, i.e., a multi-head self-attention sublayer and a position-wise feed-forward sublayer. Each block operation is defined as:
Attnl,ht = SM {
1√ da
〈 W q,l,hzlt,W k,l,hzlt′ 〉} ,
(7)
f l,tMHSA = T∑
t′=1
H∑
h=1
Attnl,ht W o,l,hW v,l,hzlt′ , (8)
f l,tFFN = W l FFN2 σ(W l FFN1 LN(f l,t MHSA + z l t)),
(9)
zl+1t = LN(f l,t FFN), (10)
where SM and LN represent the softmax and layer normalization operations, respectively. Attnl,ht is the attention score matrix between the vector zlt at position t ∈ [T ] and other vectors zlt′ at position t′ ∈ [T ] in the l-th transformer block. f l,tMHSA and f l,tFFN are the outputs of the multi-head selfattention sublayer and the point-wise feed-forward sublayer in the l-th transformer block, respectively. W q,l,h,W k,l,h,W v,l,h ∈ Rdlz×dla represent the query, key, and value weight matrices, respectively. W o,l,h ∈ Rdla×dlz represents the aggregated weights across H heads. dla is the dimension of the transformer block l, i.e., the width of the entire block. H is the number of heads and the dimension of each head in the transformer block l is dla = d l z/H . σ represents the ReLU activation function. W lFFN1 ∈ Rd l z×dlin ,W lFFN2 ∈ Rd l in×dlz represent the inner feed-forward weight matrices. dlin is the inner dimension of the feed-forward sublayer, which is usually set to 4dlz. Unlike traditional transformer networks stacking blocks with fixed sizes, in this study, we allow each transformer block l to have different dlz and d l in dimensions, enhancing its flexibility across different tasks.

heading: A.1.2 Neural Tangent Kernel (NTK)
text: Formally, assume a deep neural network f parameterized by w has D output dimensions. Let (X ,Y)
be the training samples, and L the loss function. The outputs of the network are f(X ,w) ∈ RND, where N is the number of training samples. During gradient descent training, the evolution of parameters ws and output f(X ,ws) at time step s can be expressed as follows:
ẇs= −η∇wf(X ,ws)⊺∇f(X ,ws)L, (11) ḟ(X ,ws)= ∇wf(X ,ws)ẇs
=−η∇wf(X ,ws)f(X ,ws)⊺∇f(X ,ws)L =−ηΘs(X ,X )∇f(X ,ws)L, (12)
where Θs(X ,X ) ∈ RND×ND is the Neural Tangent Kernel (NTK) at time step s, defined as:
Θs(X ,X ) = ∇wf(X ,ws)∇wf(X ,ws)⊺. (13)
The NTK exactly captures the training dynamics of the network. Especially for the infinite wide network, under the mean-squared loss and a constant NTK assumption, i.e., Θs(X ,X ) ≡ Θ0(X ,X ), Equation (12) has a closed-form solution: f (X ,ws) = ( I− e−ηΘ0s ) Y + e−ηΘ0sf (X ,w0) ,
(14)
where f (X ,ws) represents the outputs of the network at time step s, I is the identity matrix, f (X ,w0) is the output of the network at initialization, and η is the learning rate. This implies that the output of the network is determined by the training samples (X ,Y), the initial weights w0, and the initial NTK f (X ,w0). Through the NTK at initialization, we can estimate the training convergence of a network. Arora et al. (2019) have demonstrated that the training convergence speed is faster in the direction corresponding to the larger NTK eigenvalues of the network.

heading: A.1.3 Separation Rank
text: Separation rank was first proposed by Beylkin and Mohlenkamp (2002) for high-dimensional numerical analysis, and then applied to various areas including machine learning (Ghassemi et al., 2019), chemistry (Chinnamsetty et al., 2007), and physics (Validi, 2014). Recently, it has been applied to measure the input dependencies modeled by deep convolutional and recurrent networks (Cohen and Shashua, 2017; Cohen et al., 2016). Let (A,B) be a partition of the input locations, i.e., A and B are disjoint subsets of the input sequence [T ] := {1, . . . , T} and A ∪ B = [T ]. The separation rank of a function f(x1, . . . ,xT ) with respect
to partition (A,B) is the minimal number of summands that together sum up to equal f , where each summand is multiplicatively separable with respect to the partition (A,B), i.e., each summand is equal to a product of two functions — one that takes in only inputs from the subset {xi : i ∈ A} and another that intakes only inputs from the other subset {xj : j ∈ B}. Formally, the separation rank of f : ( Rdx
)T → R with respect to the partition (A,B) is defined as:
sep(A,B)(f) := min { R ∈ N ∪ {0} :
∃gA1 , . . . , gAR, gB1 , . . . , gBR : ( Rdx )N/2 → R,
f(x1, . . . ,xT ) =
R∑
r=1
gAr ({xi : i ∈ A}) gBr ({xj : j ∈ B}) } .
(15)
The separation rank quantifies the amount of input inter-dependency induced by the function f(x1, . . . ,xT ) with respect to partition (A,B). If the separation rank of a function f(x1, . . . ,xT ) with respect to partition (A,B) is 1, the function f(x1, . . . ,xT ) is multiplicatively separable with respect to partition (A,B), meaning that {xi : i ∈ A} and {xj : j ∈ B} are statistically independent. The larger sep(A,B)(f) is, the more it models inter-dependency between {xi : i ∈ A} and {xj : j ∈ B}. In other words, it means that the function f(x1, . . . ,xT ) can learn higher correlations between {xi : i ∈ A} and {xj : j ∈ B}.

heading: A.2 Proof
text: Proof of Theorem 3.2. Under the vanilla stochastic gradient descent (SGD) optimizer, the weight parameter is updated as follows:
ws+1 = ws − ηs ∂L ∂w ∣∣∣∣ w=ws
(16)
Then, under the first-order Taylor expansion, the mean of output over tokens at time step s+ 1 satisfies:
f̄s+1(X ) = f̄s(X )− ηs( ∂L ∂w ∂f̄s(X ) ∂w )
= f̄s(X )− ηsL′(f̄s)( ∂f̄s(X ) ∂w )( ∂f̄s(X ) ∂w )⊺ = f̄s(X )− ηsL′(f̄s)K̄(X ,X ). (17)
This reduces to a kernel gradient descent method. Thus, the convergence rate of the network is determined by the eigenstructure of the mean NTK K̄(X ,X ). If it can be diagonalized by eigenfunctions with corresponding eigenvalues λi, the sum of eigenvalues ∑ λi provides a upper bound for the convergence of the network. Since the mean NTK K̄(X ,X ) matrix is symmetric and symmetric. The trace of K̄(X ,X ) i.e., NTKT is equal to the sum of eigenvalues ∑ λi. Therefore, a larger NTKT score of the Transformer network indicates it converge faster.
Proof of Theorem 3.4. According to the Equation (4), the output of each block l can be expressed as follows:
f l+1t =
T∑
t′=1
H∑
h=1
〈 W q,l,hf lt ,W k,l,hf lt′ 〉 W o,l,hW v,l,hf lt′ .
(18)
Let M l,h=W q,l,hf lt f lT t W k,l,hT , the Equation (4) can be expressed as:
f l+1t =
T∑
t′=1
H∑
h=1
W o,l,hM l,hW v,l,hf lt′ . (19)
Assume there exits a balanced partition of the input locations (A,B), i.e., each token index Pt ∈ {A,B}, A and B are disjoint subsets of input sequence [T ] := {1, . . . , T} and A ∪ B = [T ]. The matrix multiplication in the M l,h can be divided the sum over inputs indexed by A and B:
M l,hr1,r2 =
T∑
t=1
[W q,l,hf lt ]r1,t[f lT t W k,l,hT ]t,r2
= ∑
t∈A [W q,l,hf lt ]r1,t[f
lT t W k,l,hT ]t,r2
+ ∑
t∈B [W q,l,hf lt ]r1,t[f
lT t W k,l,hT ]t,r2 .
(20)
Then the Equation (19) can be reformed as:
f l+1t = ∑
h∈[H]
dla∑
r1,...,rT=1
∑
P1,...,PT
W o,l,hW v,l,hf lt
( T∏
t=1
[W q,l,hf lt ]rt,t[f lT t W k,l,hT ]t,rt
)
= ∑
h∈[H]
dla∑
r1,...,rT=1
∑
P1,...,PT
W o,l,hW v,l,hf lt
(∑
t∈A [W q,l,hf lt ]rt,t[f
lT t W k,l,hT ]t,rt
∑ t∈B [W q,l,hf lt ]rt,t[f lT t W k,l,hT ]t,rt
)
≤ T∑
t=1
dla∑
r1,...,rt=1
H
=
T∑
t=1
(Hdla) t =
T∑
t=1
(dlz) t
= dlz 1− (dlz)T+1
1− dlz .
(21)
The separation rank at the l block satisfies:
log( sep(f l+1t )
sep(f lt) ) ≤ log(dlz) + log( 1− (dlz)T+1 1− dlz )
(22)
Then, the separation rank at the last block L satisfies:
log(sep(fLt ))≤ log( L∑
l=1
dlz)+log( L∑
l=1
1− (dlz)T+1 1− dlz )
(23)

heading: A.3 Algorithm Details of ETAS
text: To integrate our proposed NTSR zero-cost proxy into the ETAS framework, we first randomly sample N0 networks from the Transformer search space and compute the relative rankings of these N0 networks using NTSR. After that, we select the top n0 network structures with the highest NTSR scores as the parent population and evaluate these networks to obtain their true performance. We add the network-performance pairs to the observed set. For each iteration m, we generate a pool of Nm candidate architectures by mutating the current parent
population and select the top nm architectures from Nm based on their NTSR scores. We then evaluate these nm architectures and add their networkperformance pairs to the current observed set. The parent population is updated by selecting the top n0 networks from the current observed set. Finally, we select the top-performing architecture from the observed set. This process continues until the maximum number of iterations M is reached or the current best value has not improved for five successive iterations. The detailed algorithm of our proposed ETAS framework is summarized in Algorithm 1. Note that our proposed ETAS framework can also incorporate other zero-cost proxies.
Algorithm 1: ETAS Input: Total number of iterations M ,
search space A, parent population = ⊘, observed population = ⊘
Output: The best-performing architecture f∗
1 Randomly sample N0 networks from the search space A; 2 Compute the relative ranking of N0 initial networks by NTSR; 3 Select top n0 networks as the parent population; 4 Evaluate n0 networks and add the corresponding validation accuracy into the observed population; 5 for m = 1 to M do 6 Generate Nm candidate architectures by
mutating the network from the parent population;
7 Compute the relative ranking of Nm initial networks by NTSR; 8 Select the top nm networks and evaluate them to obtain the corresponding validation accuracy; 9 Add the nm networks and their corresponding validation accuracy into the current observed set; 10 Update the parent population by selecting the top n0 networks from the current observed set; 11 end 12 return the best-performing architecture from
the observed set.

heading: A.4 Experimental settings

heading: A.4.1 Experimental settings of Section 4.1
text: To evaluate the performance of zero-cost proxies in the ImageNet-1K dataset, we compute Spearman’s ρ (Krishnakumar et al., 2022b) and Kendall’s τ (Ning et al., 2021) rank correlation between the zero-cost proxy scores and the validation accuracy of these sampled networks. We run each experiment over 5 independent runs with different random seeds and compute the mean and standard deviation of rank correlations. As showed by Chen et al. (2021a), the initial weight inherit from the pre-trained supernet can achieve the performance comparable to that of the retrained one, we obtain the true accuracies of sampled network when they inherit their weights from the pre-trained supernet. During the search of ETAS, we first randomly sample 500 networks from the Transformer search space and compute the relative rankings of 500 networks in terms of the zero-cost proxy. After that, we choose the top-3 network structures as the parent population and evaluate these networks to obtain their true performance. For each iteration m, we generate a pool of 100 candidate architectures by mutating the current parent population and select the top-3 architectures from 100 candidate architectures in terms of the zero-cost proxy. The mutation probability is set to 0.4. The parent population is updated by selecting the top-3 networks from the current observed set. We set the number of iterations M to 10. We then find the optimal ViT network from the observed set. At last, we follow the training configuration in AutoFormer Chen et al. (2021a) to train the optimal ViT network and obtain its test accuracy on the test set of ImageNet1k dataset. All experiments are conducted in a machine with an Intel Xeon Gold 5218R CPU and two NVIDIA GeForce RTX 3090 GPUs.

heading: A.4.2 Experimental settings of Section 4.2
text: We run each experiment over 5 independent runs with different random seeds and compute the mean and standard deviation of rank correlations. We train each GPT-2 model from scratch following the settings of Radford et al. (2019). Validation perplexity is measured over a sequence length of 192 and 32 tokens for WikiText-103 and LM1B datasets, respectively. We use the BPE tokenizer and set the vocab size to 50264. For the WikiText103 dataset, we train the GPT-2 network for 4×104 steps using the LAMB (You et al., 2020) optimizer
with a batch size of 128, a learning rate of 1 × 10−2 with a cosine scheduler, and attention dropout set to 0.1. For the LM1B dataset, we train the GPT-2 network for 1× 105 steps using the LAMB optimizer with a batch size of 128, a learning rate of 3× 10−4 with a cosine scheduler, and attention dropout set to 0.1.
During the search of ETAS, we first randomly sample 300 networks from the search space to warm up the entire ETAS algorithm. We compute the relative rankings of these networks in terms of the zero-cost proxy. After that, we choose the top3 network structures as the parent population and evaluate these networks to obtain their true performance. For each iteration m, we generate a pool of 100 candidate architectures by mutating the current parent population and select the top-3 architectures from the 100 candidate architectures in terms of the zero-cost proxy. The mutation probability is set to 0.3. The parent population is updated by selecting the top-3 networks from the current observed set. We set the number of iterations M to 10. All experiments are conducted on a machine with an Intel Xeon Gold 5218R CPU and two NVIDIA GeForce RTX 3090 GPUs.

heading: Future Work
text: In the future, we hope our proposed proxy can extend to other complex types of Transformer networks like BERT (Devlin et al. , 2019) and T5 (Raffel et al. , 2020)