Neural Architecture Search as Program Transformation Exploration

Current Limitations.

NAS researchers assume the compiler is a black box bundled with the hardware, while compiler writers assume that the network architecture is set in stone. NAS researchers can discover good networks but are limited to a set of pre-implemented operations; compiler writers can efficiently exploit hardware structure but miss larger scale optimization opportunities.

NAS designs are not guaranteed to be correct and have to be separately evaluated through either a full retraining process or on a smaller proxy task. This training process severely limits the search space and can render large scale searches intractable.

What we want is the best of both worlds. We wish to combine neural architecture and compiler optimization in a unified framework. In this paper, we recast neural architecture search as program transformation exploration. By including transformations such as grouping and bottlenecking into the compiler optimization search space, we not only leverage the extensive history of program transformation research, but also discover new forms of neural architecture reduction that would not have been available to us otherwise.

Our approach.

2.1 Models and Code

The fundamental building block of a DNN is the tensor convolution, a generalization of a basic convolution.

Basic Convolution.

Consider the first row in Figure 1. This shows a basic convolution where each element in the 2D output matrix is the weighted sum of the corresponding and neighboring elements of the 2D input matrix. The weights are stored in a small 2D filter or weight matrix W whose size corresponds to the size of the neighborhood. This is shown both diagrammatically on the left and as code on the right.

2.2 Models and Code Transformations

2.3 Combined Space

If we consider swapping a full convolution for a reduced convolution as a program transformation, we can combine them. For example, we could apply loop interchange to the code in Figure 1 row 4, ($[C_o^{\prime }<C_o/B,C_i]\mapsto [C_i,C_o^{\prime }<C_o/B]$). With $C_i$ now the outermost iterator, we can reapply bottlenecking, giving input channel bottlenecking. Such an optimization is both semantically invalid, and also unavailable in existing neural architecture search spaces. However, given the robustness of neural networks to noise, in some specific cases it may be just as representationally preserving as output channel bottlenecking. It is only through the lens of program transformations over loop nests that we are able to automatically access such operators.

3.1 Neural Architectures Components

The convolution operation dominates computational complexity and is the primary component of interest in NAS.^5,18

3.2 NAS Search Space Example

Typically, the NAS heuristic finds the best cell to fill in the skeleton by choosing appropriate operations. Figure 2 gives an example cell from.⁵ The skeleton is ResNet-like with 5 cells in series. Each cell contains 4 inter-connected nodes (A,B,C,D). Between nodes, downsampling takes place where the spatial dimensions are halved in size while doubling the channel depth. This restricts the types of operations available on each edge. This gives a total of 15625 possible cells, which captures most of the available options within cell-based NAS techniques, each of which the authors train exhaustively on various datasets.

Rather than designing cells by choosing which of a predetermined list of options is best, we instead wish to choose a sequence of transformations from an original model which will allow us to step through the cell design space. We describe the program transformations we use in the next section.

4.1 Legality

Legality in the polyhedral framework is determined by preservation of data dependences. If there is a data dependence between two dynamic statement instances under the original program schedule, the relative ordering between these statements must be preserved under the new transformed schedule. Given a linear transformed schedule of the iteration space $I$ in matrix form $T$ and a dependence matrix $D$, then a transformation is legal iff:

which is to say that the lexicographical ordering of the iteration instances is preserved.

5.1 Extending the Polyhedral Model

for (g=0; g<G; g++)
  for (co=Co/G*g; co<Co/G*(g+1); co++)
    for (ci=Ci/G*g; ci<Ci/G*(g+1); ci++)
      for (oh=0; oh<OH; oh++)
        for (ow=0; w<OW; ow++)
          for (kh=0; kh<KH; kh++)
            for (kw=0; kw<KW; kw++)
              O[co][oh][ow] +=
                        W[co][ci][kh][kw] *
                        I[co][oh+kh][ow+kw];

for (g=0; g<Co; g++)
  for (oh=0; oh<OH; oh++)
    for (ow=0; w<OW; ow++)
      for (kh=0; kh<KH; kh++)
        for (kw=0; kw<KW; kw++)
          O[g][oh][ow] +=
                    W[g][kh][kw] *
                    I[g][oh+kh][ow+kw];

5.2 Fisher Potential as a Legality Check

The transformations described above fail to preserve traditional program semantics. We instead state that a schedule transformation for a neural network with respect to a task is legal if the final classification accuracy on the held out data is the same, or similar to within a small $\delta$. Training all possible transformed networks to determine accuracy, however, would be prohibitively expensive.

To address this, we employ Fisher Potential¹⁴ as a pre-training legality check. It is a cheap-to-compute measure that is effective at rejecting any architectures whose final test accuracy is significantly below the original starting point, without needing to perform training. Informally, Fisher Potential is, locally, the total information that each loop nest (layer) contains about class labels under a simplifying assumption of conditional independence.

Example.

Let us consider, candidate networks designed from NAS-Bench-201⁵ as shown in Figure 3. Each point represents a different neural architecture of which there are 15625. The $y$-axis shows final CIFAR-10 top-1 error (lower is better), and the $x$-axis shows the Fisher Potential assigned to the model at initialization (higher is better). We can see that without requiring any training, Fisher Potential is able to filter out poorly-performing architectures, visible in the cluster of low scoring networks on the left with poor final errors. Many good networks are also discarded — this is unfortunate but acceptable for our scenario, since the space is large and densely populated with networks.

5.3 Expressive Power

Neural Architecture Search techniques rely on hand-designed exploration spaces. themselves are engineered by hand. For example, a recent paper found that bottlenecking could be applied in the spatial domain.¹¹ This can now be added to a list of candidate operations. However, spatial bottlenecking is automatically captured within our framework as the combination of existing transformations. As an example, consider the convolution in row 2 of Figure 1. We use the shorthand notation $\left[i\right]\stackrel{B\left(b\right)}{\to }\left[i\left(b\right)\right]$ to denote bottlenecking domain $i$ by factor $b$ and $[i,j]\stackrel{\text{int}}{\to }[j,i]$ to denote interchange. Then the spatial bottleneck operation can be constructed as the following sequence of transformations:

The expressiveness of the polyhedral model, equipped with these new transformations, means that there are novel transformations that can be derived from this unified framework.

7.1 CIFAR-10 Network Performance

The results for each network are shown in Figure 4. The combination of TVM and NAS is a strong baseline, however, for each of the models our unified method is able to generate models with improved hardware performance. In general there is more performance to gain on the mGPU than other platforms, as relaxed memory pressure from smaller designs is of increasing importance.

ResNet-34: NAS is able to find a modest improvement of 1.12$\times$ speedup over TVM on the GPU platform, increasing to 2$\times$ on i7 CPU, showing the power of compressed convolutions. Our unified approach is able to find further improvement giving 2$\times$ and 3$\times$ speedup respectively. The improvement is more dramatic on the smaller platforms with 5$\times$ and 10$\times$ speedups on the mCPU and mGPU.

ResNext-29-2x64d: NAS is unable to find any improvement here due to the already highly compact structure of the network. There are simply no NAS options that improve performance over the TVM baseline across all platforms. Despite this our approach is able to find small improvements of 1.3$\times$ and 1.1$\times$ on the CPU and GPU platforms by combining NAS and program transformations. This increases to 1.4$\times$ on the embedded CPU and 7$\times$ on the mGPU.

DenseNet-161: The impact of NAS is much more varied here. While it can find 2.2$\times$ improvement on the CPU, it has a negligible improvement over TVM on the GPU. It also struggles on the mCPU but finds 6$\times$ on the mGPU. Our approach is able to improve by over 3$\times$ for both server platforms, but only finds 1.2$\times$ on the mCPU while acheiving 10$\times$ on the mGPU.

7.2 Analysis

Accuracy. For all CIFAR-10 networks pictured in Figure 4, changes in accuracy were less than 1% in absolute difference. The ResNet-34 in Figure 5 had an original ImageNet Top-1 and Top-5 accuracy of 73.2% and 91.4% respectively. The final network compiled with our transformations had a Top-1 and Top-5 accuracy of 73.4% and 91.4%; it was slightly more accurate than the original but considerably faster. To give a broader overview, we present the accuracy of several ImageNet models in Figure 7, which shows that for each model accuracy degradation is small or non-existent.

Size: One benefit of these neural transformations is their effect on model size, both in terms of weights and runtime memory usage. We found that CIFAR-10 networks could be compressed in size by $2–3\times$. Likewise, the ImageNet ResNet-34 in Figure 5 was compressed from 22M parameters to 9M without a loss in accuracy.

Search time: Our search process involves suggesting configurations and rejecting or accepting them based on Fisher Potential. Since Fisher Potential is extremely cheap to compute, the search time overhead introduced by our method is small, less than 5 minutes on a CPU. During this time we were able to explore 1000 different configurations, discarding approximately 90% of the candidate transformation sequences through the Fisher Potential legality check.

7.3 Transformation Sequence Case Studies

There were 3 particular transformation sequences that dominated the list of best performing transformations. In neural architecture terms, the resulting operations of each of these sequences of transformations is a new convolution-like operator, previously unavailable to NAS.

 Sequence 1 applies grouping to the kernels over the spatial domain of the input. These are then concatenated to form one output. The exact sequence is: [split $\to$ interchange $\to$ group $\to$ interchange $\to$ fuse].

 Sequence 2 is an operator in which the output channels have been unrolled by factor 16 and then the remaining iteration domain has been grouped by factor $G=2$. The exact sequence is: [unroll $\to$ group $\to$ interchange]  Sequence 3 is an operator that was devised by splitting up the iteration domain of the output channels, and applying different levels of channel grouping to each new domain (e.g. $G=2$ on first half, $G=4$ on the second half). The exact sequence is: [split $\to$ group $\to$ interchange $\to$ group]. Each of them is applicable across all networks and may be more widely applicable as standard tensor convolutions for other networks.

7.4 Exploring Layer-Wise Optimizations

In Figure 5, we examine the impact of our sequences layer-by-layer on one network on one platform: ResNet-34 on the i7. The exact configurations mirror the experiment in the original TVM paper.³ While we achieve a 3x improvement overall as shown in Figure 4, the speedup achieved varies considerably across layers. In 4 of the 11 layers no performance improvement is found, as Fisher Potential marks these individual layers to be extremely sensitive to compression. Simple grouping with a factor $G=2$ is able to give around 2$\times$ speedup across 7 of the 11 layers. Sequence 1 is able to give a small improvement in most cases — particularly layer 11 — where spatial reduction yields new parallelization opportunities. Sequence 2 provides a slight improvement in many cases, and through the later layers it is able to offer larger speedups through improved data reuse.  Sequence 3 is the best option for most early layers but suffers compared to other sequences in the later layers.

7.5 Alternative NAS comparison

To provide an alternative evaluation against an existing NAS technique, we re-implement FBNet¹⁶ using the convolutional blocks available in our NAS space, and our three baseline networks as the skeletons into which the selected cells are inserted. Figure 6, shows the performance of resulting networks on CIFAR-10 relative to TVM, NAS and our approach on the Intel Core i7 for each network. In each case, FBNet does provide a modest improvement over NAS. However, it is worth noting that this involves an expensive training step at each stage of evaluation (requiring $\sim$3 GPU days per network). Our approach is able to consistently improve over FBNet, with no training required.

7.6 ImageNet Network Performance

We also applied our technique to the ImageNet classification dataset, to show that the method transfers beyond CIFAR-10. In Figure 7, we first show the resultant networks for running our method on ResNet-18, ResNet-34, DenseNet-161, DenseNet-201, and DenseNet-169, with their inference times recorded on the Intel i7 CPU. For each network, accuracy is within 2%, and there are significant gains in inference time.

7.7 Interpolating Between Models

The additional expressive power of our method allows us to explore the search space of networks in a more fine-grained manner than traditional NAS techniques. This is because NAS techniques simply choose operations from a list, whereas we generate new ones. To illustrate this consider Figure 8 where we plot the accuracy and inference time of a ResNet-34 with two BlockSwap-based models (the blue points labeled NAS-A and NAS-B). We can explore alternatives by a series of parametrized transformations in our framework, yielding the points in red. Each of these points is a new possible blocktype that would not be accessible to a traditional NAS technique unless explicitly written by the human designer. During this interpolation we are also able to find a Pareto optimal point.