SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Yun Yi $^{1,3,\dagger}$  Haokui Zhang $^{2,3,*}$  Rong Xiao $^{3}$  Nannan Wang $^{1,*}$  Xiaoyu Wang $^{3}$ $^{1}$ Xidian University  $^{2}$ Northwestern Polytechnical University  $^{3}$ Intellifusion yuny220@163.com hkzhang1991@mail.nwpu.edu.cn  
xiao.rong@intellif.com nnwang@xidian.edu.cn fanghuaxue@gmail.com
As more deep learning models are being applied in real-world applications, there is a growing need for modeling and learning the representations of neural networks themselves. An effective representation can be used to predict target attributes of networks without the need for actual training and deployment procedures, facilitating efficient network design and deployment. Recently, inspired by the success of Transformer, some Transformer-based representation learning frameworks have been proposed and achieved promising performance in handling cell-structured models. However, graph neural network (GNN) based approaches still dominate the field of learning representation for the entire network. In this paper, we revisit the Transformer and compare it with GNN to analyze their different architectural characteristics. We then propose a modified Transformer-based universal neural network representation learning model NAR-Former V2. It can learn efficient representations from both cell-structured networks and entire networks. Specifically, we first take the network as a graph and design a straightforward tokenizer to encode the network into a sequence. Then, we incorporate the inductive representation learning capability of GNN into Transformer, enabling Transformer to generalize better when encountering unseen architecture. Additionally, we introduce a series of simple yet effective modifications to enhance the ability of the Transformer in learning representation from graph structures. In encoding entire networks and then predicting the latency, our proposed method surpasses the GNN-based method NNLP by a significant margin on the NNLQP dataset. Furthermore, regarding accuracy prediction on the cell-structured NASBench101 and NASBench201 datasets, our method achieves highly comparable performance to other state-of-the-art methods. The code is available at https://github.com/yuny220/NAR-Former-V2.
With the maturity of deep learning technology, an increasing number of deep neural network models of various sizes and structures are being proposed and implemented in academic research and industrial applications. In this process, the rapid deployment of networks and the design of new networks that meet task requirements are significant. To address this issue, researchers propose using machine learning models to solve the deployment and design problems of the models themselves. One popular strategy is encoding the input neural network and utilizing the resulting neural network representation to predict a specific target attribute directly without actually executing the evaluation program. In recent years, we have witnessed success in accelerating model deployment and design
This work was done while Yun Yi was an intern at Intellifusion.
processes with the help of neural network representations [47, 21, 3, 46, 26, 44, 12]. Taking the advantages of latency predictors [47, 21, 17, 55, 12, 2], significant time cost and expertise efforts can be saved by not having to carry out the time-consuming process of compilation, deployment, inference, and latency evaluation when engineers choose networks for application. Through the use of accuracy predictors [47, 26, 46, 4, 12, 44, 20], researchers can avoid the resource-intensive process of network training and instead perform a forward inference process to evaluate the accuracy of a multitude of networks. This measure dramatically reduces the time cost associated with network design.
Although the vanilla Transformer is designed for natural language processing, Transformer architecture has found widespread adoption across diverse fields owing to its strengths in global modeling and parallelizable computation [24, 22, 23, 14, 41, 8]. Very recently, several researchers have attempted to learn appropriate representations for neural networks via Transformer [47, 26]. These methods have indeed achieved leading performance on relevant tasks. Nevertheless, they are mainly designed for encoding the architecture of cells (basic micro units of repeatable neural networks) in cell-structured networks. As shown in the latency prediction experiment in NAR-Former (Transformer-based neural architecture representation learning framework) [47], poor generalization performance occurs when the depth of the input architecture reaches hundreds of layers. In the development process of neural network representation learning, Graph neural network (GNN) [16, 35] is also a promising technique for learning neural network representations [21, 37, 19, 43, 12]. They model the input neural architecture as a directed acyclic graph (DAG) and operate on the graph-structured data, which comprises the node information matrix and adjacency matrix. Recently, the NNLP [21] introduced a dedicated latency prediction model based on GNNs, which is capable of encoding the complete neural network having hundreds of layers and achieves a cutting-edge advance.
In fact, both cell-structured architectures and complete neural networks are widely used in various applications. Cell-structured models offer good scalability, allowing for easy scaling by adding or removing cells. This adaptability makes them suitable for addressing problems of different complexities and data sizes, while also facilitating incremental model development and deployment. Complete neural networks provide better flexibility in connectivity and can achieve higher accuracy in certain cases. Furthermore, in some cases, such as in latency estimation, encoding the complete network is necessary. To handle various network architectures in different tasks, both GNN-based and Transformer-based models are necessary. However, this issue of utilizing multiple architectures can introduce constraints that may not be conducive to practical applications. For instance, when a designed network requires specific attributes, having similar model structures and high-accuracy predictions for different attributes can reduce code redundancy and improve work efficiency.
In this paper, we build upon the research conducted in NAR-Former [47] and present a novel framework called NAR-Former V2 for universal neural network representation learning. Our framework can handle cell-structured networks as well as learn representations for the entire network. To accomplish this, we incorporate graph-specific properties into the vanilla Transformer and introduce a graph-aided attention-based Transformer block. This approach combines the strengths of both Transformer and graph neural networks (GNNs). Extensive experiments are conducted to evaluate our proposed framework. Results show that: (1) our method can be applied to predict different attributes, can outperform the state-of-the-art method in latency prediction on the NNLQP dataset [21], and can achieve promising results in accuracy prediction on the NAS-Bench-101 and NAS-Bench-201 datasets [48, 10]; (2) our method has good scalability, which is capable of encoding network having only a few operations or complete neural networks that have hundreds of operations.
2.1 Representation and attribute prediction of neural networks
Neural network representation learning is the base for evaluating the attributes of different networks via machine learning models. Early methods [7, 20] construct representation models for learning neural network representation based on LSTM and MLP. Peephole [7] inputs the embedding of each layer to LSTM to predict accuracy, which neglects the topological structure and is limited to handling only sequential architectures. Later, in order to better capture the structural information of the network, an accuracy predictor [44] uses a binary path encoding with a length equal to the number of possible paths from input to output given in terms of operations, where the element at the corre
sponding position of the path presenting in the input network is set to 1. When the neural network is regarded as a directed acyclic graph, the adjacency matrix describes the connection between nodes, so it is naturally used to encode the topological structure of the neural network. NAS-Bench-101 [48] proposed to encode the given neural network as a concatenated vector of a flat adjacency matrix and a list of node labels. Many other methods [21, 17, 3, 12, 43, 19] realize accuracy and latency prediction by directly inputting the original two-dimensional adjacency matrix together with the node information matrix to GNN, which can realize the explicit encoding of the input network topology. There also are some works exploring other ways of obtaining representations of neural networks, such as by mimicking actual data processing [31] or message exchange [50] of neural networks, and by deriving a number of persistent topology measures for DNNs [5]. Recently, other methods have focused on enhancing the original GNN [4] or introducing transformers [26, 47] to obtain more meaningful neural network representations.
2.2 Transformer
Transformer [41] is a self-attention-based neural network architecture that has revolutionized natural language processing [28, 8, 33] and has been adopted in many other fields [54, 24, 22, 23, 47, 26, 11, 14]. Because of its excellent modeling capability and wide range of applications, there is a growing interest in the direction of further enhancing its performance [53, 1, 39, 6], improving computing efficiency [23, 40, 34], combining with other backbone networks [15, 18, 38, 29, 13], and even taking inspiration from it to optimize other types of network structures, such as ConvNets [45, 51, 25, 9, 52]. Transformer has recently been successfully introduced into neural network representation learning [26, 47]. TNASP [26] inputs the sum of the operation type embedding matrix and Laplacian matrix into the standard Transformer. NAR-Former [47], on the other hand, encodes each operation and connection information of this operation into a token and inputs all tokens into a proposed multi-stage fusion transformer. Excellent attribute prediction results have been achieved on cell-based datasets by using these methods. However, the strong long-range modeling ability of the self-attention mechanism may also result in subtle local variation affecting the representation of all tokens. Due to the potential impact of this feature on the generalization ability, although NAR-Former [47] has made attempts to encode complete neural networks, the results are still unsatisfactory.
2.3 Graph neural network
GNNs are designed to handle graph-structured data, which is a fundamental representation for many real-world problems such as social network analysis and recommendation systems [49, 32]. Given that neural networks can be viewed as graphs, GNN-based models have emerged as a prominent and widely adopted approach for neural network representation learning [21, 17, 3, 12, 43, 19]. GNNs show generalization ability through a simple mechanism of aggregating information from neighbors. For instance, the recently proposed GNN-based model [21] can obtain representations of neural networks with hundreds of layers and achieves new state-of-the-art results in latency prediction, even if the input network structure has not been seen during training. Nevertheless, the simple structural characteristics of GNNs, which contribute to their strong generalization ability, also lead to the need for further improvement in the performance of methods based on original GNN in cellular structure and complete neural network representation learning. Therefore, it is a promising approach for neural network representation learning to combine the Transformer and GNN to leverage the strengths of both models.
Our final framework diagram for realizing neural architecture representation learning and attribute prediction is shown in Fig. 1, consisting of three main phases: neural network encoding, backbone-based representation learning, and attribute predicting using prediction heads. We will introduce the motivation and details of the backbone network design in Sec. 3.1 and Sec. 3.2, and the details of the entire model in Sec. 3.3.
Figure 1: Overview of attribute prediction model.
As mentioned in the Sec. 1, Transformer-based models have demonstrated remarkable performance in encoding and learning representations of neural networks when the input is in the form of cells. However, when dealing with complete deep neural networks (DNNs) consisting of hundreds of layers, or when the depth of the input data is unseen during training, they may sometimes exhibit poorer performance compared to GNN-based methods. Additionally, as highlighted in [21], real-world applications often show significant differences in the topologies and depths between training and test samples. Consequently, representation learning models must possess strong generalization abilities for handling unseen data. In this regard, GNN-based models appear to achieve better performance.
This observation has prompted us to reconsider the two types of inputs, namely cells and complete DNNs, as well as the two representation learning models, the Transformer and GNN. Through a detailed comparative analysis of the structures of the Transformer and GNN, we speculate that the insufficient generalization capability of Transformer-based methods may be attributed to its structure and computation characteristics. As we know, the self-attention structure in transformers is a crucial design that allows for the effective extraction of global features in a data-driven manner. However, this structure becomes a double-edged sword when learning network representations. For input neural networks with depths of hundreds of layers, the Transformer's impressive capability to capture global information can sometimes lead to excessive sensitivity. This stems from the fact that the Transformer models interactions between all tokens using its self-attention mechanism, treating the entire sequence as a fully connected graph. This dense attention mechanism can give rise to a particular issue: even a subtle variation, such as reducing the kernel size in a layer from  $5 \times 5$  to  $3 \times 3$ , can affect the representation of all other layers, ultimately leading to significant differences in the final representation. As a result of this issue, the trained model may be biased toward fitting the training data distribution, it yields inferior results and demonstrates poorer generalization performance. The corresponding experiments are presented in Sec. 4.4.
3.2 Transformer grafted with GNN
Fig. 2 shows the vanilla transformer block, GNN block, and our proposed graph-aided attention Transformer block. As shown in Fig.1 (a), the vanilla Transformer block has two major parts:
$$
\hat {H} ^ {l} = \operatorname {S e l f A t t n} \left(\mathrm {L N} \left(H ^ {l - 1}\right)\right) + H ^ {l - 1}, \tag {1}
$$
$$
H ^ {l} = \operatorname {F F N} \left(\operatorname {L N} \left(\hat {H} ^ {l}\right)\right) + \hat {H} ^ {l}, \tag {2}
$$
where  $H^l$  is the feature for the layer  $l$ .  $\hat{H}^l$  is an intermediate result. SelfAttn, FFN, and LN refer to self-attention, feed-forward network, and layer normalization, respectively. GNN block just has one major part, where the representation is updated following:
$$
\hat {H} ^ {l} = \operatorname {G r a p h A g g r e} \left(H ^ {l - 1}, A\right) + W _ {r} ^ {l} H ^ {l - 1}, \tag {3}
$$
$$
H ^ {l} = \mathrm {L} _ {2} (\hat {H} ^ {l}), \tag {4}
$$
where  $\operatorname{GraphAggre}(H^{l-1}, A) = W_a^l(\operatorname{Norm}(A)H^{l-1})$ .  $A \in \mathbb{R}^{N \times N}$  is the adjacency matrix, and  $\mathrm{L}_2$  denotes the 12-normalization function.  $W$  with different superscripts and subscripts represents different learnable transformation matrices.
(c)
Figure 2: Diagrams of three modules. (a) The vanilla Transformer block [42]. (b) The GNN layer with mean aggregator [16]. (c) The proposed graph-aided attention Transformer block.
- The Transformer utilizes self-attention to fuse information from a global perspective, while the GNN uses graph aggregation to fuse neighbor information based on the adjacency matrix.
- The Transformer block includes an additional FFN (Feed-Forward Network) component, which enhances information interaction between channels.
The advantage of self-attention lies in its data-driven structure, allowing for flexible adjustment of information fusion weights based on the input data. On the other hand, the advantage of GNN is that graph aggregation focuses on the topological structure. These two advantages are not contradictory to each other. Consequently, we have naturally come up with an idea to combine the strengths of self-attention and graph aggregation. This approach inherits the flexibility of self-attention while benefiting from the good generalization capability of graph aggregation.
To implement this idea, we consider the neural network encoded as a graph, with the operations or layers in the network treated as nodes. Assuming the graph has  $N$  nodes, the transformer layer we have designed for universal neural network representation learning (Fig. 2 (c)) is calculated as follows:
$$
\widetilde {H} ^ {l} = \operatorname {T A E n h a n c e} \left(H ^ {l - 1}, D\right), \tag {5}
$$
$$
\hat {H} ^ {l} = \mathrm {L} _ {2} \left(\operatorname {G r a p h A t t n} \left(\widetilde {H} ^ {l}, A\right) + W _ {r} ^ {l} \widetilde {H} ^ {l}\right), \tag {6}
$$
$$
H ^ {l} = \operatorname {G F F N} \left(\operatorname {L N} \left(\hat {H} ^ {l}\right)\right) + \hat {H} ^ {l}, \tag {7}
$$
where  $D \in \mathbb{R}^{N \times 1}$  is a vector that records the number of nodes directly connected to each node. The Graph-aided Attention (GraphAttn) module is responsible for performing attention calculations using the properties of the graph structure to adjust global self-attention. The Type-aware Enhancement module (TAEnhance) is utilized to further enhance the representation. We introduce a Grouped Feed-Forward Network (GFFN) by introducing group linear transformation into the original FFN. In the following sections, we will provide a detailed introduction to each component.
Graph-aided attention In the proposed graph-aided attention, we employ the adjacency matrix to govern the attention calculation range. Moreover, the adjacency matrix characterizes the interlayer connection relationships within the neural network, enabling the model to acquire topology knowledge. Hence, we define this module as the Graph-aided Attention module:
$$
X ^ {l} = \operatorname {S i g m o i d} \left(W _ {q} ^ {l} \widetilde {H} ^ {l} + b _ {q} ^ {l}\right), \tag {8}
$$
$$
S ^ {l} = \left(X ^ {l} X ^ {l T} / \sqrt {d}\right) \odot A, \tag {9}
$$
$$
Z ^ {l} = W _ {a} ^ {l} \left(\operatorname {N o r m} \left(S ^ {l}\right) \widetilde {H} ^ {l}\right) + b _ {a} ^ {l}. \tag {10}
$$
The  $\mathrm{Norm}(\cdot)$  means that for each node, the attention weights between it and other nodes are transformed to  $(0,1)$  by dividing it by the sum. The character  $b$  with different superscripts and subscripts represents different learnable biases. The  $d$  refers to the feature dimension of  $X^l$ . Note that simply using the adjacency matrix to control the attention map in self-attention is insufficient. To make this approach work, we have discarded the original softmax operation in self-attention and replaced it with a linear attention mechanism. This is because the softmax operation tends to focus excessively on the current node while neglecting neighboring nodes. Consequently, we have inserted a sigmoid activation function before the linear attention to ensure that all values in the attention map are positive. For further comparisons with the original self-attention, please refer to the supplementary.
Type-Aware enhancement module The connection between a layer and other layers is related to the type of that layer. Therefore, the number of connected layers in each layer can be used to assist the model in learning the type of layer. By fully utilizing the internal characteristics of this graph-structured data, it is beneficial to improve the learned representations. The enhanced representation is obtained by:
$$
\operatorname {T A E n h a n c e} \left(H ^ {l - 1}, D\right) = \operatorname {S i g m o i d} \left(W _ {d} ^ {l} D + b _ {d} ^ {l}\right) \odot H ^ {l - 1}. \tag {11}
$$
3.3 Universal representation learning for neural network
In this subsection, we will present the construction of a comprehensive framework for neural network encoding and representation learning based on the proposed enhanced Transformer block. We will also explain how this framework can be utilized to predict network attributes. The overall system, as depicted in Fig. 1, is composed of three consecutive stages: neural network encoding, representation learning, and attribute prediction.
Neural network encoding We have taken inspiration from the tokenizer used in NAR-Former [47] and made certain modifications to encode the input neural network. For a given network consisting of  $N$  layers or operations, whether it is a cell architecture or a complete DNN, we represent it as a sequence feature comprising vectors corresponding to each layer:  $T = (t_{1}, t_{2}, \dots, t_{N}) \in \mathbb{R}^{N \times C}$ . In NAR-Former, for each vector of the current node, it encapsulates information of operation, self-position, and the position of its source nodes:  $t_{i} = (t_{i}^{\mathrm{op}}, t_{i}^{\mathrm{self}}, t_{i}^{\mathrm{sour}}) \in \mathbb{R}^{C}$ .
Following the encoding scheme proposed in NAR-Former [47], we use the position encoding formula [41, 30] to transform the single real-valued numbers (e.g. operation labels and node position indices) of relevant information into a higher-dimensional space. We denote this mapping scheme as  $f_{\mathrm{PE}}(\cdot)$ .
For the operation encoding  $t_i^{\mathrm{op}}$ , there are slight differences in the specific encoding content for input networks of different scales. If the input is in the form of a cell architecture [48, 10], there are usually no more than ten different options for each architecture operation. In this case, we directly assign category labels to all candidate operations, and then use the function  $f_{\mathrm{PE}}(\cdot)$  to encode the label of the operation to obtain  $t_i^{\mathrm{op}}$ . However, when a complete DNN is used as input [21], more abundant operational information can be extracted and encoded. In this case, we first use one-hot vectors, which ensure the same distance between different categories, to encode the type of operation (e.g. convolution, batch normalization, ReLU, concatenation). Then use the function  $f_{\mathrm{PE}}(\cdot)$  to encode the properties (e.g. kernel size, number of groups) of the operation, which is then concatenated with the one-hot type vector as  $t_i^{\mathrm{op}}$ . NAR-Former [47] uses the concatenation of  $t_i^{\mathrm{self}}$  and  $t_i^{\mathrm{sour}}$  to encode information flow. Since our improved transformer can obtain the topology information of the network with the help of the adjacency matrix, these two parts of the encoding are not necessarily needed in our method. Details about the encoding in the different experiments are provided in our released code and supplementary material.
Representation learning The model for learning neural network representations  $H^{K}$  is constructed by stacking multiple instances of our proposed enhanced Transformer blocks. These improvements, specifically tailored for neural network data, allow the model to learn representations that are more meaningful and exhibit enhanced generalization capabilities.
Attributes predicting Taking the representation  $H^K$  as input, the target attribute can be predicted by using the predicting head:
$$
\hat {y} = - \text {l o g s i g m o i d} (\mathrm {F C} (\operatorname {R e L U} (\operatorname {F C} (\operatorname {R e L U} (\operatorname {F C} (H ^ {K})))))). \tag {12}
$$