SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Mohammad Salameh $^{1*}$ , Keith G. Mills $^{1,2*†}$ , Negar Hassanpour $^{1}$ , Fred X. Han $^{1}$ , Shuting Zhang $^{3}$ , Wei Lu $^{1}$ , Shangling Jui $^{3}$ , Chunhua Zhou $^{3}$ , Fengyu Sun $^{3}$ , Di Niu $^{2}$
<sup>1</sup>Huawei Technologies Canada. <sup>2</sup>Dept. ECE, University of Alberta. <sup>3</sup>Huawei Kirin Solution, China. {mohammad.salameh, negar.hassanpour2, fred.xuefei.han1, jui.shangling, zhouchunhua}@huawei.com,
{kgmills, dniu}@ualberta.ca, {zhangshuting8, robin.liuwei, sunfengyu}@hisilicon.com
Optimizing Deep Neural Networks (DNNs) to obtain high-quality models for efficient real-world deployment has posed multi-faceted challenges to machine learning engineers. Existing methods either search for neural architectures in heuristic design spaces or apply low-level adjustments to computation primitives to improve inference efficiency on hardware. We present Automated Graph Optimization (AutoGO), a framework to evolve neural networks in a low-level Computation Graph (CG) of primitive operations to improve both its performance and hardware friendliness. Through a tokenization scheme, AutoGO performs variable-sized segment mutations, making both primitive changes and larger-grained changes to CGs. We introduce our segmentation and mutation algorithms, efficient frequent segment mining technique, as well as a pretrained context-aware predictor to estimate the impact of segment replacements. Extensive experimental results show that AutoGO can automatically evolve several typical large convolutional networks to achieve significant task performance improvement and FLOPs reduction on a range of CV tasks, ranging from Classification, Semantic Segmentation, Human Pose Estimation, to Super Resolution, yet without introducing any newer primitive operations. We also demonstrate the lightweight deployment results of AutoGO-optimized super-resolution and denoising U-Nets on a cycle simulator for a Neural Processing Unit (NPU), achieving PSNR improvement and latency/power reduction simultaneously. Code available at https://github.com/Ascend-Research/AutoGO.
Deep Neural Networks (DNNs) have achieved great success in Computer Vision (CV) and Natural Language Processing (NLP) tasks. A major trend toward achieving better performance on benchmarks is adopting large and computationally demanding deep models [7]. However, successful and efficient deployment of deep neural networks onto diverse and specific hardware devices, including neural processing units on the edge, significantly hinges upon engineering proper neural architectures that are both excellent in task performance while meeting hardware friendliness objectives.
A range of techniques have been proposed by academia and industry to solve the hardware-friendly deployment challenges of DNNs [46]. Neural Architecture Search (NAS) replaces the manual design process of DNNs, achieving remarkable performance in several applications in CV [67, 11, 9, 6] and NLP [32, 10, 12]. While NAS can utilize a flexible range of search algorithms [53, 11, 79, 47], the search space adopted by NAS is based on heuristics, either searching for an optimal macro-net
†Work done during an internship at Huawei.
construction based on predefined blocks, e.g., MBConv blocks in MobileNets [59, 26], or stacking searchable cells by heuristic rules [39, 17, 73]. These heuristic rules may not be efficient to the target hardware device for deployment and may limit the potential gain from NAS methods. On the other hand, graph rewriting [70, 30] operates on the tensor computation graph of a DNN to improve its inference efficiency on hardware. Rewriting involves applying a set of subgraph substitution rules that preserve mathematical functionality of the original DNN, which does not alter or reduce the neural architecture to achieve better task performance or fitness to hardware.
In this paper, we propose Automated Graph Optimization (AutoGO), a generic graph optimization framework to evolve a given neural architecture for efficient and low-power deployment onto a specific hardware device. Unlike traditional NAS which builds networks from scratch in a heuristic search space or from hand-crafted building blocks [39, 63, 15], AutoGO enhances both hardware-friendliness and task performance of a DNN, by evolving its underlying Computation Graph (CG) using computational units composed of operations extracted from NAS benchmarks. AutoGO automatically improves typical neural networks in terms of benchmark performance on several CV tasks ranging from classification, semantic segmentation, to super-resolution without relying on newer operations. It also automates lightweight DNN deployment onto mobile neural processing units while preserving task performance, hence replacing manual tweaking efforts by ML engineers. We present the following key contributions in designing the AutoGO framework:
Figure 1: A DNN can be partitioned into disjoint subgraphs (segments). Segments contain a variable number of inputs, outputs, and nodes ranging from primitive operations to complex subgraphs.
AutoGO optimizes DNNs on a Computation Graph (CG) of low-level primitives extracted from TensorFlow [1] models, which allow us to optimize all types of operations and their hyperparameters like filter size and latent tensor dimensions and thus offer a holistic fine-grained view of DNNs. Unlike graph rewriting which preserves mathematical equivalence, AutoGO alters the CG of a DNN for task performance and fitness to hardware.
Rather than relying on predefined blocks, the basic units for mutation in AutoGO are computation subgraphs, which we call segments, as illustrated in Figure 1. Segments are mined from a large number of CGs from several NAS bench
marks based on frequent subgraph mining in a data-driven manner. We use Byte Pair Encoding (BPE), an efficient tokenization technique [20] from NLP, to segment CGs and merge frequent operations into segments. By extracting and including segments of variable sizes into our database, AutoGO enables both primitive operation changes and larger-scaled block changes to a DNN.
AutoGO leverages an evolutionary algorithm to mutate our BPE-segmented source network. The segment mutation process is guided by hardware friendliness metrics and a pre-trained neural predictor to estimate the performance of the mutant network resulting from segment replacement. We propose a neural predictor which explicitly considers the positional and contextual information of a segment replacement made in a CG and directly models the performance gain. Mutations are also coupled with a resolution propagation scheme that solves for the tensor shapes in the replacement segment to ensure architectural validity.
Extensive experiments demonstrate that AutoGO can enhance the performance of the best architectures in several public architecture benchmarks, e.g., NAS-Bench-101 [71], NAS-Bench-201 [17], HiAML, Inception, and Two-Path [48]. Additionally, AutoGO can automatically optimize several typical large CNN architectures, including ResNets [24], VGG [61], and EDSR [38] on a breadth of CV tasks including Classification, Semantic Segmentation, Human Pose Estimation, and Super Resolution. We show that AutoGO can improve their performance while making them lightweight, without using newer generations of operations that do not appear in the original network. Finally, to demonstrate the real-world applicability of our framework, we show results of AutoGO-optimized FSRCNN [16] (for super-resolution) and image denoising U-Net [55] for low-power or low-latency deployment using a cycle simulator for a Huawei mobile Neural Processing Unit (NPU).
NAS Benchmarks and Neural Predictors. NAS-Benchmarks comprise architectures from a given search space and their accuracy performance. NAS-Bench-101 [71] and NAS-Bench-201 [17] provide the performance of 423k and 15.6k architectures, respectively, on CIFAR-10. Benchmarks enable the rapid development of search algorithms and neural predictors. Neural predictors [43, 54, 41, 73, 65, 40] treat NAS benchmarks as datasets and learn to estimate the performance of architectures in a given search space, and thus constitute a low-cost avenue for performance evaluation.
However, NAS benchmarks suffer from several drawbacks. First, benchmarks only provide performance annotations for architectures inside a manually designed fixed search space. Thus, any tweaks to decrease FLOPs or latency beyond the search space requires training the new architecture from scratch. Second, existing NAS Benchmarks are mostly cell-based [39] and compose a network by stacking the same cell structure multiple times. This structure forms a high-level architecture representation that is insensitive to spatial details such as latent tensor dimensions, which vary along the depth of the network and influence hardware-friendliness [49]. As most existing neural predictors learn using high-level cell representations, these drawbacks hamper their deployment generalizability. In contrast, AutoGO can mutate an architecture beyond its original, manually-defined design space by utilizing a low-level, spatially-sensitive representation.
Computational Graphs for DNN Hardware Friendliness. Multiple subfields explore how to reduce the carbon footprint and time cost of DNNs. Pruning and quantization methods [36] aim to reduce the number of parameters or lower the bit precision of model weights, respectively. Graph rewriting methods like TASO [30] and TENSAT [70] consider mathematically equivalent substitutions, e.g., merging or splitting parallel convolutions and applying the associative/distributive properties. These schemes usually require spatial details like resolution and channel size to perform rewrites. Hence, they represent DNNs using Computation Graphs (CG) [50, 22], which treat each primitive operation as a node and use the network forward pass to define edge connectivity. Similarly, AutoGO also uses a low-level CG representation. But different from these approaches, it aims to evolve the architecture of an untrained DNN to improve performance while also optimizing hardware friendliness.
Neural Architecture Design Space. Several works employ NAS over large spaces by jointly searching over macro and micro-structures for both block type and tensor dimensions [64, 15]. Human expertise is at the core of these design choices to constrain the search for high-performing architectures. [57] model a search space as a 3-level hierarchical graph to overcome the reliance on expert knowledge. [13] propose Neural Search Space Evolution, which progressively grows a current search space by adding unseen operation candidates. Unlike the above work, we do not limit ourselves to a pre-designed skeleton with spatial or topological constraints at any network position. Rather, we incorporate search space and architectural knowledge into a neural predictor. Also, instead of manually defining the search units [17, 59, 62, 42], we mine these units from NAS benchmarks. In particular, we utilize Frequent Subgraph Mining (FSM) [31] to discover interesting and frequent patterns in the computation graphs in a data-driven way. FSM requires conducting expensive steps when graphs are large such as extracting patterns, inspecting isomorphism, and checking if subgraphs are frequent enough to be considered interesting. Algorithms [68, 69, 27] trade result completeness and accuracy for efficiency to overcome run time and memory inefficiency [19]. In NAS, [56] utilize Weisfeiler-Lehman (WL) graph kernel to extract useful network features but only applies it shallow cell-based DAG structure of NAS-Bench-201. In contrast, we propose an efficient approach to mine frequent subgraphs by converting CGs to sequences and applying BPE [20] to extract subsequences, which produces a fixed-size vocabulary of subgraphs of varying sizes.
3 AutoGO: Automated Computation Graph Optimization
The AutoGO framework operates on the Computation Graph (CG) of an input DNN architecture extracted from the in-memory graph structure of a tensorflow.keras model or .pb file. CGs are directed acyclic graphs (DAGs), where nodes represent primitive operations that are indecomposable computation operations in deep learning frameworks like ONNX [5] and PyTorch [52], e.g., Convolutions, Pooling, ReLU, Add, etc., while edges represent forward-passes between operations. Specifically, node features include operation type, input and output tensor resolution dimensions and weight tensor shape if applicable.
Figure 2: AutoGO takes the CG of a neural network as input and improves it using an algorithm-mined segment database and pre-trained mutation performance predictor.
Figure 2 provides an overview of the proposed AutoGO operating on the CG level. AutoGO mutates the CG of an input DNN by utilizing a database of segments (Sec. 3.1), which are frequent subgraphs mined from a variety of NAS-Benchmarks via an efficient tokenization method. A Pareto front evolution strategy (Sec. 3.2) performs segment mutations to the CG, while using resolution propagation to ensure network validity. Finally, a pretrained Predecessor-Segment-suCcessor (PSC) neural predictor (Sec. 3.3) together with selected hardware metrics will guide the evolution by assessing the performance gain when a certain segment is stitched into the architecture.
3.1 Computation Graph Segmentation via Tokenization
We partition a CG  $g$  into a contiguous sequence of subgraphs, which we call segments. A segment, denoted by  $s$ , may have multiple input and output nodes, or could even be a single primitive operation. For  $g$  to be a valid neural network, any two contiguous segments  $s_i$  and  $s_j$  in  $g$  must maintain the correctly matched height, width and channel  $(H, W, C)$  resolutions. The resolutions of output nodes of  $s_i$ , have to match the resolutions of input nodes of its succeeding segment  $s_j$ . Our definition of segment spans a wider range of topologies and provides flexible operation grouping than predefined or handcrafted blocks, e.g., ResNet and MBCnv blocks or DARTS cells [39].
The Segment Database  $\mathcal{D}$  is the core component that stores the segment units that AutoGO mutations are based on. AutoGO uses the segment database  $\mathcal{D}$  (or vocabulary) to either partition an input CG into segment units or select a segment from  $\mathcal{D}$  to replace an existing segment in input CG. To alleviate the memory and time complexities of mining common subgraphs from a large number of CGs, we relax the problem into mining segments from sequences. This allows us to utilize much more efficient tokenization techniques over sequences.
Given a set of neural networks represented in a CG format  $\mathcal{G} = \{g_1, \dots, g_k\}$ , we enrich node representation by labeling a node in the form of [current op, incoming ops, outgoing ops]. Each unique node label is further encoded into a single character symbol. Thus, a topologically sorted CG can be mapped into a string of character symbols, where each character encodes an operation and its neighboring operations. By converting all graphs in  $\mathcal{G}$  into sequences, we essentially have built a corpus for training a tokenizer to extract common segments of character symbols. Specifically, we use Byte-Pair Encoding (BPE), a data compression [20] technique with a wide use for text tokenization in NLP [60], to tokenize the string representations of CGs. BPE operates iteratively by collecting frequent pairs of consecutive symbols to build a vocabulary of tokens (segments). Using BPE, we extract the most common subsequences from the string representation of the CGs and build a vocabulary of size  $|V|$ . We revert each discovered subsequence in the vocabulary back to its corresponding subgraph representation from the CG to form a segment database  $\mathcal{D}$ . Given a new CG, BPE utilizes its built vocabulary and applies a greedy algorithm to segment it. Figure 1 provides an example of a segmented CG. Our approach brings several benefits over mining on large graphs with WL-kernels. The extraction process on sequences is efficient. Using BPE enables segment extraction from all benchmark families simultaneously without facing memory inefficiencies like WL-Kernel.
For example, we label a Conv 3x3 node with incoming edges from Add and BatchNorm operations and an outgoing edge to a ReLU operation as "conv2d3,in,add,batchnorm,out,relu".
3.2 Computation Graph Mutation
We use an evolutionary search strategy to perform segment mutations on the CG of an input architecture and iteratively update a Pareto front of architectures in terms of predicted accuracy and a chosen hardware-friendliness objective, e.g., FLOPs, latency, and power. The mutation made to a parent architecture comprises the following steps: segmentation, source segment selection, replacement segment selection, tensor resolution (shape) propagation, and performance evaluation. First, AutoGO partitions the parent architecture into segments using the BPE-generated vocabulary  $V$ . BPE adopts a greedy segmentation approach, which will lead to deterministic partitioning. To diversify the segmentation outcome, we select a subset of vocabulary  $V' \subset V$  that BPE uses during segmentation, thus leading to different partitioning outcomes every time the CG is segmented. After segmentation, we select a set of candidate source segments to mutate. For each source segment  $s_i$ , we randomly select multiple replacement segments  $s_i^* \in \mathcal{D}$  that have the same number of inputs and outputs as the corresponding source segment. If  $s_i^*$  has multiple inputs and/or outputs, AutoGO randomly maps these to the outputs and/or inputs created by removing  $s_i$ .
The mutation process must maintain a valid architecture, by correctly combining the replacement segment with the rest of the model. Given a CG  $g$ , let  $S_g = \{s_0, s_1, \ldots, s_{n-1}\}$  be the set of disjoint segment subgraphs generated by applying BPE to  $g$ . Let  $s_i, 0 \leq i \leq n-1$  be a source segment within  $g$  that we wish to replace. We partition the segments within  $S_g$  into three distinct groups that reflect their positions within  $g$ :
- The Predecessor group denotes all segments  $P = \{s_p; 0 \leq p < i\}$  between the input and  $s_i$ .
- The specific Segment  $s_i \in \mathcal{D}$  that we are aiming to replace by mutating  $g$ .
- The suCessor group denotes all the remaining network segments  $C = \{s_c; i < c \leq n - 1\}$ .
Let  $\{P, S, C\}$  refer to a CG partitioned in this manner, denoted by grey, purple and orange blocks in Fig. 3. Hence, for a mutation to be valid, the shape of the output tensor from the Predecessor  $P$  must match that of the input to the replacement segment  $s_i^*$  and the output shape of  $s_i^*$  must match the shape of the expected input to suCdecessor  $C$ . AutoGO adapts replacement segment  $s_i^*$  to the remainder of the architecture  $P$  and  $C$  by adjusting the hyperparameters of operations in  $s_i^*$  to achieve the desired resolutions. Adjustments are applied to mutable operations, e.g., increasing the stride of convolutions and pooling operations to induce downsampling, whereas operations such as activations and batch normalization are immutable. Depending on the  $P, C$  and  $s_i^*$  subgraphs, stitching the replacement segment into the overall CG may be infeasible. We cast this "resolution propagation" task as a Mixed-Integer Linear Program (MILP) over the
adjustable hyperparameters of mutable nodes within  $s_i^*$ , which is an optimization problem with linear objectives and constraints, and integer-valued decision variables. We define MILP constraints that regulate the correct resolution propagation within  $s_i^*$  when stitched to the rest of the architecture.
Figure 3: An example of segment mutation. AutoGO takes an input CG (left) and replaces a source segment  $s_i$  with  $s_i^*$ . The predecessor (grey) and successor (yellow) are unchanged.
At the end of each mutation iteration, we retain all the segment replacements that have a feasible solution to the resolution propagation MILP, and profile these valid segment mutations in terms of the predicted accuracy gain given by the PSC predictor (Sec. 3.3) and a selected hardware friendliness metric, based on which a Pareto front  $\mathcal{O}$  [21] of architectures is maintained and updated. During the first iteration, we only mutate the input architecture. At the beginning of each successive iteration, AutoGO selects the  $k$  architectures from the Pareto front as parents to mutate. If the Pareto front contains less than  $k$  architectures, AutoGO will select additional non-Pareto optimal architectures having the minimum sum of accuracy and FLOPs ranks. Further technical details on the  $\{P,S,C\}$  partitioning scheme, resolution propagation, parent selection process and overall algorithm, including examples, illustrations and pseudocode, are provided in supplementary Sections A.3.1 and A.5.
3.3 Context-Aware Mutation Performance Estimation
AutoGO uses a neural predictor to estimate the performance of mutant architectures in order to search for high-quality networks. We propose a novel predictor that assesses the potential performance benefit from a segment mutation, based on its topology, context and its location within the architecture.
Just as we can construct subgraphs from the individual segments produced by BPE, we can construct larger predecessor and successor subgraphs from all segments within  $P$  and  $C$ , respectively. We use this format to train a PSC neural predictor. Let  $h_*$  denote a fixed-length graph embedding for an arbitrary graph produced by a graph neural network (GNN) [51]. The PSC predictor operates by separately encoding the predecessor, segment, and successor CGs into distinct graph embeddings. These embeddings are then concatenated and fed into a multi-layer perceptron (MLP) regressor to predict the performance  $y$  of  $g$ . Stated more formally:
$$
h _ {P} = \operatorname {G N N} (P); h _ {s _ {i}} = \operatorname {G N N} \left(s _ {i}\right); h _ {C} = \operatorname {G N N} (C); y = \operatorname {M L P} \left(\operatorname {C o n c a t} \left[ h _ {P}, h _ {s _ {i}}, h _ {C} \right]\right).
$$
It is also possible to predict performance using the entire CG,  $y = \mathrm{MLP}(\mathrm{GNN}(g))$ . However, our PSC predictor enjoys several advantages over this approach. By separately encoding the  $\{P,S,C\}$  partitions, the PSC predictor is sensitive to the position and context of the segment  $s_i$  within the overall network. This allows us to more directly estimate the performance impact that mutating  $s_i$  will have. As shown in Figure 3, we are considering mutating  $g$  into  $g^*$  by replacing the source segment  $s_i$  with a segment  $s_i^*$ . We want the mutant network to outperform the original, i.e.,  $y^* > y$ . Our training process emphasizes learning a separate embedding for each  $s_i$  in our segment database  $\mathcal{D}$ . It encodes the required knowledge to estimate the effect of small changes from segment  $s_i$  to any replacement segment  $s_i^*$ , given a fixed  $P$  and  $C$ .
We construct our database by extracting segments from 5 CIFAR-10 [33] benchmark families: NAS-Bench-101 (NB-101) [71], NAS-Bench-201 (NB-201) [17], HiAML, Inception, and Two-Path [48]. Initially, we set the BPE vocabulary size to 2000 and obtain a database with 428 unique segments after filtering out isomorphisms. Segments vary in size ranging from primitives (containing a single operation node) to blocks with up to 16 nodes and edges. The average segment contains 5 nodes and 3 edges, and some segments are disconnected subgraphs spanning parallel branches of a CG. We provide in-depth statistics and visualizations in Section A.1.