CAP: A Context-Aware Neural Predictor for NAS

1 Introduction

Neural architecture search (NAS) [Elsken et al., 2019] has gained increasing attention for its capability to automatically design well-performing deep neural networks in various realworld applications [Sun et al., 2019; Ghiasi et al., 2019; Liu et al., 2019]. However, NAS often needs to evaluate derived architectures with the help of numerous computational resources, which are often not affordable for most researchers. To tackle this issue, various methods are proposed to reduce the search cost, such as early stopping [Ru et al.,

∗Corresponding author.

2021], using proxy dataset [Zela et al., 2020], weight sharing [Pham et al., 2018], and so on. Although these methods indeed improve the efficiency of NAS, they suffer severely from performance degradation, making the performance of architectures searched cannot be guaranteed in practice [Xie et al., 2023].

Neural predictor is a promising acceleration method which can directly estimate the performance of an unseen architecture accurately [Wen et al., 2020]. In practice, it requires extensive architecture-performance pairs to learn such mapping on the specific tasks. However, these neural predictors tend to perform well only when enough reliable annotated architectures are available, but obtaining such architectures requires a repetitive training process, which is still time-consuming.

Several works have tried to train a well-performing predictor when a small amount of annotated architectures are available. For example, Semi-NAS [Luo et al., 2020] attempts to generate more annotated architectures in a semi-supervised way and HAAP [Liu et al., 2021b] proposes a homogeneous architecture augmentation strategy for predictors to expand the training data. However, neither focuses on improving architecture representation to enhance the predictor. We observe that the massive unlabeled architectures in the search space can potentially help the predictor produce expressive representations of architectures but are rarely explored at present. If these data are leveraged properly for training the predictor, it is unnecessary to annotate a large amount of architectures. Namely, the search cost for predictor-based NAS can be greatly reduced.

In view of the observation above, we propose a contextaware neural predictor (CAP) and design a context-aware self-supervised task for neural predictors to make full use of the substantial unlabeled architectures. By predicting the contextual structure around subgraphs in each architecture, the predictor is enabled to understand and capture its intrinsic topology information and node features. Specifically, we extract the central subgraph and the corresponding context graphs in every graph-like architecture, to construct graph pairs. During this pre-training stage, the graph data from the same pair can be embedded into similar representations. Following the commonly used fine-tuning process in selfsupervised learning [Liu et al., 2021a], we apply a simple regression model to map the architecture representation to the real performance and then utilize a few annotated architectures to fine-tune the predictor. Consequently, CAP can capture the topological information at the structure level and the node features at the attribute level of architectures.

Our contributions can be summarized as follows:

• We propose a context-aware neural predictor (CAP) which requires much fewer annotated architectures for training. Compared with the existing neural predictors, CAP can still rank architectures in the search space precisely even with fewer annotated architectures.
• We design a context-aware self-supervised task to pretrain the neural predictor, which is capable of generating meaningful and generalizable representations of architectures without involving annotated architectures. This helps predict the performance of architectures precisely and search for well-preforming architectures in NAS.
• The experimental results show that CAP can evaluate the performance ranking of architectures precisely in NAS-Bench-101 and NAS-Bench-201 search spaces with fewer annotated architectures. Specifically, the number of annotated architectures needed by CAP is 2.4× fewer than that of most existing predictors as shown in Figure 1. Moreover, CAP can achieve satisfactory search results in NAS-Bench-101, NAS-Bench201, and DARTS search spaces. In addition, the ablation studies demonstrate the effectiveness of the proposed context-aware self-supervised task for neural predictors.

2 RelatedWorks

2.1 Neural Predictor

Neural predictor is an effective method to evaluate the performance of architectures and consequently guides the optimal searching direction in NAS. Generally, most of the existing predictors have three stages: 1) Sufficient architectures are collected and then trained from scratch; 2) a predictor is constructed after utilizing these annotated architectures to learn their mapping to the real performance; 3) the predictor can be applied to estimate unseen architectures rapidly and accurately. In these approaches, only a small portion of labeled architectures are selected to construct the predictor due to the unacceptable computation cost required for training. Several works concentrate on making full use of the limited annotated architectures from the perspective of information flow [Ning et al., 2020], operation hierarchy [Chen et al., 2021b], and spatial topology information [Lu et al., 2021], respectively. In contrary to the aforementioned methods, we dive into the substantial but less explored unlabeled architectures. After pre-training the predictor on these highly accessible data to learn their intrinsic features, very limited annotated architectures are enough for a well-performing predictor.

We note that some previous works try to utilize unlabeled architectures as well, which is similar to our work. Arch2Vec [Yan et al., 2020] uses the vanilla graph reconstruction pretext task to obtain representations from architectures in an unsupervised manner. Instead of focusing on the vanilla graph reconstruction which simply ignores the distinctiveness of the architectures during encoding, we propose a contextaware pretext task to make full use of the rich contextual information from massive unlabeled architectures. To extract node attributes among architectures, GMAE-NAS [Jing et al., 2022] constructs a difficult graph reconstruction pretext task by masking a large portion of nodes in the architectures to pre-train the predictors. However, GMAE-NAS only focuses on the attribute of operations in the architectures. In contrast, the CAP method further considers the structural information by introducing a context-aware pretext task, enriching the predictor with more valuable prior knowledge.

2.2 Context-Aware Self-Supervised Learning

Self-supervised learning is gaining mounting attention for its successful applications in various domains. Typically, it defines a pretext task based on unlabeled input data to obtain descriptive representation for different downstream tasks. Context-aware self-supervised learning is a significant branch, which constructs pretext tasks mainly according to contextual information inside the data. In natural language processing, a typical context-aware task is to predict the contextual sentences of a given word. In this way, the model is expected to explore the relationship of words in the text without involving any label. Similarly in computer vision, [Doersch et al., 2015] predicts the relative positions of contextual patches in an image and enables the model to produce significant feature representation. Moreover, TS-TCC [Eldele et al., 2021] also learns time-series representation from unlabeled data by contextual contrasting. In light of the successful applications above, we propose a context-aware self-supervised task to predict the contextual surrounding in an architecture. The predictor is encouraged to obtain intelligible representations of architectures and thus relies on fewer annotated samples.

3 Methodology

As shown in Figure 2, the overall framework of the proposed CAP method contains three parts, i.e., the architecture encoding, the context-aware self-supervised task, and the performance prediction.

3.1 Architecture Encoding

Following the conventions of predictor-based NAS algorithms, a neural architecture can be viewed as a standard graph where the nodes denote the various operation layers and the edges denote the various connection manners. Thus, a N-layer architecture X is naturally encoded as an operation list ON = {o1, o2, ..., oN} and an adjacent matrix AN×N . In the operation list ON, oi represents the operation type of the i-th layer.

However, not all search spaces treat the operations of architectures as nodes in the graph. To bridge the gap, we shift the form of the operation-on-edge search spaces (e.g., the NAS-Bench-201 [Dong and Yang, 2019a] and DARTS [Liu et al., 2018] search spaces) to the standard version via a simple transformation.

Considering that an architecture could be represented as different isomorphic graphs, we choose a variant version of Graph Isomorphism Networks (GINs) [Xu et al., 2018] as the encoder to obtain invariant and expressive representation from given architectures. Generally, the update mechanism of the k-th layer of a GIN can be formulated as follows:

where hv and hu denote the features of the node v and u, Hθ is a simple MLP, N(v) indicates the set of all nodes within the neighborhood of node v, and ϵ is a learnable weight.

3.2 Context-Aware Self-Supervised Task

Once the architectures are encoded into graph data, we conduct the context-aware self-supervised task to pre-train the predictor. As for graph self-supervised learning, there exist two kinds of popular pretext tasks: predictive pretext tasks and contrastive pretext tasks. The former aim to predict the self-generated labels to provide self-supervision for the model. The latter produce various augmentation views for a given graph sample and then maximize the similarity of those matching pairs generated from the same sample. We select the latter to leverage the rich contextual information among architectures because contrastive learning methods are on the ground of mutual information maximization and predictive learning methods lack a unified theoretical framework at present [Xie et al., 2022]. Inspired by the application of successful context-aware self-supervised tasks in natural language processing [Mikolov et al., 2013], computer vision [Doersch et al., 2015] and times-series [Eldele et al., 2021], we allow the encoder to predict the contextual surrounding of any node in each architecture. The basic idea is to first extract the central subgraph and the context graphs from each architecture to construct graph pairs. After that, the similarity between matching pairs which are obtained from the same architectures is maximized.

In the context-aware self-supervised task, we randomly choose a central node v inside the graph which denotes an architecture. Then, we exploit its K-hop neighborhood GKv to represent the selected node, containing all the nodes and edges at most K hops away from node v. Once given the central node v and its neighborhood GKv , the surrounding context graphs are confirmed accordingly. We introduce a distance scale R (R > K) as the higher bound of context graphs. To be specific, all subgraphs between the K-hop neighborhood and R-hop neighborhood of the central node are regarded as the context graphs. Notably, some nodes are shared between the K-hop neighborhood and the outer context graphs, which are called anchor nodes, functioning as a bridge to convey information from the central node v to its context graphs.

For context-aware abilities, the encoders learn to match the neighborhood representation to the corresponding context representation inside an architecture by maximizing their similarity. Therefore, the K-hop neighborhood and its outer context graphs are first embedded into hGv and cGv by the encoders, respectively. Similar to other context-aware tasks, these context graphs are aggregated to represent the contextual information. Therefore, an auxiliary GIN is necessary to aggregate representation from context graphs of different sizes next to the neighborhood and generate a representative embedding cGv for the similarity computation.

Following the common practice in contrastive learning, we introduce the negative sampling strategy to train a discriminating GIN encoder. In the proposed context-aware selfsupervised task, we expect the encoder to predict the contextual structure based on the central subgraph. Hence, we randomly choose the embeddings of central subgraphs from other graphs as negative samples. We take the (hGv , cG′ v ) pairs from the same graphs as the positive ones (G = G′ ) and those from different graphs as the negative ones (G ̸= G′ ). To be specific, the negative sampling ratio is set to one, achieving the balance between positive pairs and negative ones. The loss function of the context-aware self-supervised task can be formulated as:

where (hGv , cGv ) is a positive pair from the same graph G while (hGv , cG′ v ) composes a negative pair; LCE denotes the cross entropy loss function.

After that, a readout function is added in the tail to obtain the graph-level representation of every architecture X:

The readout function here can be various graph-level pooling operations, we use averaging here for its simplicity.

The pre-trained main GIN model can be directly used as the encoder for the neural predictor. Additionally, the auxiliary GIN model is only used in the pre-training procedure and discarded in the following performance prediction.

3.3 Performance Prediction

After the pre-training process, a regression model is required to evaluate the real performance based on the representations of corresponding architectures. For simplicity, we choose a two-layer MLP as the regression model.

To fine-tune the pre-trained predictor with limited annotated architectures, we select decoder-only fine-tuning, full fine-tuning, and partial fine-tuning methods.

Decoder-only Fine-tuning. All weights of parameters in the encoder are fixed once pre-training is done. During the fine-tuning, only the regression model is trained while the encoder is frozen. In other words, the representation of any given architecture remains unchanged.

Full Fine-tuning. The whole pre-trained encoder and the regression model are jointly trained with a few labeled architectures. In this way, the pre-training task helps to initialize the encoder with better weights in comparison to random initialization.

Partial Fine-tuning. During the fine-tuning stage, partial fine-tuning discards the dropout layers and batch normalization layers in the pre-trained encoder. Thus, it takes less time to fine-tune the predictor compared with the full fine-tuning method.

As for the loss function, mean squared error (MSE) and ranking loss are two popular choices in previous works of neural predictors. Intuitively, MSE between the absolute performance of an architecture and the predicted value is effective. In this way, the model is expected to predict the real performance precisely. Unfortunately, a low MSE loss may mislead the predictors. MSE loss is unable to reflect the correct ranking between the queried architectures when their accuracy is close [Xu et al., 2021]. Conversely, it is more important to focus on the relative performance ranking of architectures considering that the goal is to search for the best architecture instead of its absolute performance in the NAS methods. Here, we use Bayesian personalized ranking (BPR) [Rendle et al., 2012] loss for the sake of its popularity in relative ranking tasks. The loss function can be formulated

where D represents all the architecture pairs like (Ai,Aj) in the training set, s(Ai) is the predicted performance score of the architecture Ai, and σ denotes the sigmoid function.

4 Experiments

The experiments performed contain three parts in terms of the ranking capability, the search results, and the ablation studies. In the first part, we verify the powerful ranking capability of CAP on NAS-Bench-101 [Ying et al., 2019] and NAS-Bench-201 [Dong and Yang, 2019a] benchmarks, especially when the training samples are extremely limited. In the second part, we use both closed domain search spaces (NAS-Bench-101, NAS-Bench-201) and open domain search space (DARTS [Liu et al., 2018]) to search for promising architectures on CIFAR-10, CIFAR-100 [Krizhevsky, 2009], and ImageNet16-120 [Chrabaszcz et al., 2017]. In the third part, we further investigate the effectiveness of the proposed context-aware self-supervised task via detailed ablation studies.

4.1 Experimental Settings

NAS-Bench-101 Search Space. The NAS-Bench-101 is a cell-based search space consisting of over 423k unique neural architectures. In each architecture cell, at most seven nodes and nine edges are contained. The nodes represent operation layers and the edges represent connection manners. The search space provides the validation accuracy and test accuracy of convolutional neural networks (CNNs) on CIFAR-10.

NAS-Bench-201 Search Space. The NAS-Bench-201 is also a cell-based search space constructed by over 15K different neural architectures. Each cell includes four nodes and six edges. Different from the operation-on-node NAS-Bench101 search space, the operations are represented by edges in NAS-Bench-201. Detailed results of CNNs are available on three datasets: CIFAR-10, CIFAR-100, and ImageNet16-120.

DARTS Search Space. The DARTS search space is built by both normal cells and reduction cells where each cell is composed of seven nodes and 14 edges. It is a much larger open domain search space than the NAS-Bench-101 and the NAS-Bench-201 search spaces. In the DARTS search space, the operations are also represented by edges.

4.2 Ranking Results on NAS Benchmarks

The experiments are carried out on both NAS-Bench-101 and NAS-Bench-201 benchmarks. In parallel to previous works [Liu et al., 2021b; Lu et al., 2021], this work utilizes the validation accuracy of an architecture to train the predictor and evaluates its corresponding test accuracy. In terms of metrics for ranking ability, we choose Kendall's Tau between the predicted accuracy of architectures and the real accuracy for its widespread use in prior works. For fair comparisons among different neural predictors, we adopt the same data splits in TNASP [Lu et al., 2021] for both NAS-Bench-101 and NAS-Bench-201.

The ranking results on NAS-Bench-101 are shown in Table 1. It can be seen that CAP outperforms the other predictors across various data splits. Notably, with only 172 annotated architectures, CAP is on par with the competitors which are trained by 424 training samples and can even beat most of them. Such promising results imply that the contextual information from unlabeled architectures is leveraged in the pretraining stage. Therefore, the pre-trained encoder can produce meaningful representations of architectures even before the fine-tuning stage. As a result, fine-tuning the predictor is much easier compared with training it from scratch. With the pre-trained predictor, only a few annotated architectures are enough to construct a well-performed neural predictor. The satisfactory ranking results demonstrate the effectiveness of the proposed context-aware self-supervised task.

Moreover, we report the ranking results on NAS-Bench201 in Table 2. It can be observed that CAP achieves stateof-the-art performance when compared with other competitors regardless of the data split. It is noteworthy that CAP surpasses them by a large margin, especially in the condition of limited training samples, which is consistent with our previous analysis of the effect of the context-aware selfsupervised task. Meanwhile, the proposed CAP can achieve a Kendall's Tau value of 0.684 when trained by only 156 annotated architectures, which is superior to other predictors trained by 469 (2.4× more than 56) training samples. The experimental results show that the proposed context-aware selfsupervised task helps construct a powerful neural predictor with fewer annotated architectures. We observe that as the number of training samples increases, the gap between CAP and other predictors is gradually narrowed. However, we mainly focus on the performance of the predictor in the fewshot scenario because annotating architectures is extremely time-consuming.

4.3 Searching on Closed Domain Search Spaces

We combine the proposed CAP with the regularized evolutionary algorithm [Real et al., 2019] to search for promising CNNs in the NAS-Bench-101 benchmark. The total number of queried architectures is set to 150 and the search results are shown in Table 3. The traditional methods are placed in the first row, and the second row shows the search results of previous state-of-the-art methods. Compared with them, CAP helps search for better architectures with the same queried budget. Through 10 independent runs, we also find that the CNNs searched by the CAP method achieve 94.18% average test accuracy on CIFAR-10, which is 0.38% higher than the result of the original regularized evolutionary algorithm. These promising experimental results demonstrate that the pre-trained predictor, which serves as the optimal guidance, facilitates searching for promising architectures.

On NAS-Bench-201, we follow the settings in [Dong and Yang, 2019a] for a fair comparison with the baselines. Specifically, 50 annotated architectures are randomly chosen to train the predictor, and the predictor afterwards evaluates all the architectures in the search space. The best performance of the top 50 evaluated architectures is reported as the search result. As demonstrated in Table 4, our method can complete the searching within 16 seconds, consuming over 80× less than the time of baselines. At the same time, the best architecture we search for exceeds the competitors in terms of accuracy. These extraordinary results can be attributed to the sufficient prior knowledge that the predictor learns in advance through the context-aware self-supervised task. Meanwhile, the applied searching strategy leverages the ranking capability of the predictor and thus dramatically cuts the searching cost.

The evaluation strategies during the search can be roughly divided into conventional training-based strategies like reinforcement learning and training-free strategies like zero-cost proxies. It can be observed from Table 4 that the former tend to discover architectures with stable performance on different independent runs at the price of enormous searching costs. While the latter are sensitive to varied settings and therefore result in a huge discrepancy of search results across the runs. Compared With them, the CAP method achieves the most balanced results on CIFAR-10, CIFAR-100, and ImageNet16120 and the searched architectures achieve state-of-the-art performance, further verifying the positive effects of the proposed context-aware self-supervised task.

4.4 Searching on Open Domain Search Space

To search for well-performed CNNs in the DARTS search space, we follow the main procedures in CATE [Yan et al., 2021]. Despite the massive available architecture performances in the NAS-Bench-301, their accuracy can be inaccurate for they are predicted by a surrogate model trained on 60k architectures. Consequently, we choose to train the architectures from scratch to obtain reliable annotated data for the predictor. Specifically, we randomly choose 100 architectures in the DARTS and then train them from scratch to construct the architecture-performance pairs as training samples. For each candidate architecture, the epoch number is set to 50 with a batch size of 96. As for the searching method, we sample 100k architectures at random in the DARTS search space and they are afterwards evaluated by our trained predictor. After that, we select the architectures with top-3 predicted performance as the search results and re-train them with common DARTS strategies. The search results on CIFAR-10 are reported in Table 5.

For a fair comparison, we report the average test accuracy of the top-3 architectures along with the best accuracy. The first row shows the performance of well-known NAS methods, and predictor-based NAS methods are placed in the second row. It can be observed that CAP obtains the best accuracy of 97.58% and the average accuracy of 97.46%, outperforming all the others with the least model parameters of 3.3M. More detailed experimental settings for the proposed pretext task and the performance evaluation are reported in the supplementary materials.

4.5 Ablation Study

We conduct ablation studies on the proposed CAP in terms of different hyperparameters using the NAS-Bench-101 search space. S = {S1, S2, S3, S4, S5} denotes the data splits in Table1. The Kendall's Tau values of predicted accuracy of architectures are reported. We also compare the searching results of CAP with/without pre-training in the supplementary materials.

Fine-tuning Methods. We train the predictor from scratch to compare it with the pre-trained versions. Furthermore, we explore two different fine-tuning methods for the pre-trained predictor, and the results are shown in Figure 3(a). It's observed that the pre-trained predictor is greatly enhanced when the encoder and the regression model are jointly fine-tuned (e.g., partial fine-tuning and full fine-tuning method). This can be credited to the rich contextual information introduced from architectures through the context-aware self-supervised task. The decoder-only fine-tuning method performs poorly because the encoder keeps frozen and lacks supervision from annotated architectures when fine-tuning, whose results are included in the supplementary materials. We also discover that the partial fine-tuning method achieves better performance than the full fine-tuning method because of a faster convergence rate.

Loss Functions. We adopt MSE loss and BPR loss for training the predictor in the fine-tuning stage. According to Figure 3(b), the predictor trained by BPR loss achieves better ranking results across various training data splits. We attribute this to its emphasis on the relative ranking of architectures rather than their absolute performance valves.

Number of Encoder Layers. We investigate the number of layers L in the GIN encoder for the predictor and report the ranking results in Figure 3(c). The optimal number L is three and deeper GIN leads to the decrease in predictor performance.

Hyperparameters in Pre-Training Task. In the proposed context-aware self-supervised task, the choice of K and R also has effects on the ranking results of the predictor. They are summarized in Figure 3(d). We find that the ranking performance is satisfactory when the central subgraph in each architecture and the corresponding context graphs share only a few overlapping nodes (i.e. K = 1, R = 2).

4.6 Visualization of Architecture Representation

To demonstrate the effectiveness of the proposed contextaware self-supervised task directly, we visualize the representations of architectures generated by the pre-trained CAP and the baseline predictor using the t-SNE tool [Van der Maaten and Hinton, 2008] in Figure 4. For the pre-trained CAP, architectures with close performance are embedded into nearby representations and the predictor could therefore rank architectures precisely using only a few annotated architectures. Instead, the baseline predictor fails to group architectures with similar performance well. This discrepancy is consistent with the results in the last section where the pre-trained CAP outperforms the baseline predictor. It also reflects that the proposed context-aware self-supervised task has positive impacts on enhancing the neural predictor.

5 Conclusion

In this paper, we propose CAP with a context-aware selfsupervised task, which leverages the rich contextual information among unlabeled architectures. The pre-trained predictor can produce expressive and generalizable representations of architectures, thus requiring fewer annotated samples. Through extensive experiments conducted on NASBench-101, NAS-Bench-201 and DARTS search spaces, the proposed CAP outperforms peer neural predictors in terms of both relative ranking and searching for well-performing architectures. The ablation studies further verify the effectiveness of the context-aware self-supervised task. CAP offers a promising direction for the predictor-based NAS methods. In the future, we will explore more self-supervised tasks for neural predictors on search spaces composed of larger units instead of cell-based ones.

Acknowledgments

This work was supported by National Natural Science Foundation of China under Grant 62276175.