HW-Aware Initialization of DNN Auto-Tuning to Improve Exploration Time and Robustness

Combining search-space knobs forms a regular, cartesian grid. When interpreting this grid as a graph, it is possible to analyse the spatial distribution of valid and invalid configurations in the search spaces. To interpret the search space as a graph, two nodes are connected if the Manhattan distance [10] between their indices in the grid is exactly one.

Only drawing edges connected to a valid candidate yields the graph Figure in 1. This graph shows a clustering of valid configurations, indicating a spatial relationship in the search space. This behaviour is not unique, as it occurred in all search spaces of Table II. This observation is valuable, as it can be leveraged to improve the optimization process. However, a detailed analysis of the effects forming clusters is not deemed helpful, as every workload configuration creates invalid configurations for many reasons.

As AutoTVM adds invalid configurations with a penalty to the performance model training set, it poses the question, how well the model can handle this. Specifically, if valid but low performing and invalid configurations are related in the search space and if treating them similarly distorts the prediction ability of the model. The following sections present the used methods to quantify the effects of invalid configurations on the ranking performance.

After evaluating the performance model in isolation, this section takes a look at the full auto-tuning process. The experiments use an untrained performance model, as is the default in AutoTVM. Tuning starts with a randomly sampled batch of size $e=50$ and a total tuning duration of $t=750$ trials. The optimization was repeated 20 times for all workloads in Table II.

The performance of the auto-tuner is defined as the ability to find good configurations fast and consistently. The following experiments use the median and interquartile range (IQR) as a measure for the former and latter criterion, respectively. The median gives a good indication on the convergence time of the auto-tuner. IQR on the other hand describes central half of all median values in multiple experiments and is useful to determine how robust the overall process is.

Figure 3 reports the amount of trials necessary to find the best performing configuration, aggregated over all 20 runs for each layer. The first, and most important observation is the wide IQR across most workloads.

For most layers, the tuner’s convergence time varies greatly, hinting at large inconsistencies in the performance model’s quality. Figure 4 shows all runs for layers 59, 92 and 107 which offers additional information on the behaviour of individual runs. For example, in layer 59 and 92 some runs did not find the best configuration during the 750 trials. This is shown by the traces of individual runs in Figure 4, that do not reach the upper performance bound. It also shows that the initial set of 50 randomly selected configurations greatly influence how long it takes the tuner to converge towards a good solution. Bad initial configurations can lead to plateaus in the search, which can last over multiple epochs. It is suspected, that the improvements found after several hundred trials can be attributed to random components in the search, rather than model quality.

Instead of randomly selecting the initial measurement batch, a balanced set of valid and invalid configurations is used. Validity information is obtained by a presampling algorithm, exploring the search space and storing information about the validity of each encountered configuration. By leveraging the spatial locality of valid candidates found in Section III-A, the randomness is reduced and more valid samples are retrieved from the search space. To check validity, the existing VTA compiler is used as a checking function.

Algorithm 1 describes how this feature is used to find more valid samples in the search space. In a first step, a randomly sampled set of configurations $P_{0}$ from search space $S$ is classified as valid or invalid (line $3$). Each configuration $p\in P$ is added to output set $N$, together with the validity information. If $p$ is valid, locality is exploited by adding all its neighbours to the set of potential candidates for the next evaluation step $C$ (line $10$). Two configurations are neighbours, if their Manhattan distance is 1. The set of neighbours has a higher chance of containing additional valid configurations. To maintain exploration, each invalid configuration adds a random new candidate from $S$ to $C$ (lines $11-12$). After processing all members of $P$, a new set $P$ is sampled from $C$, as shown in line 13. These steps are repeated until the limit of $n_{samples}$ is reached. The valid and invalid candidates to generate the balanced initialization set for AutoTVM are sampled from $N$.

Presampling, as described in the previous section, generates a set $N$ of configurations and information on their validity. From $N$, the initial measurement epoch $E_{0}\subset N$ is selected, such that a desired ratio of valid and invalid configurations is achieved. To do this, $N$ should to be larger than $E_{0}$, to ensure a degree of diversity in the possible candidates for $E_{0}$. A result of presampling algorithm is that $N$ can be defined as $N=N^{valid}\cup N^{invalid}$. Therefore, $E_{0}$ can be constructed from the configurations $E_{0}^{valid}\subset N^{valid}$ and $E_{0}^{invalid}\subset N^{invalid}$ as $E_{0}=E_{0}^{valid}\cup E_{0}^{invalid}$. The selection from $E_{0}^{valid}$ and $E_{0}^{invalid}$ happens individually, maximizing the Manhattan distance between all selected candidates.

Since the validity of a search space subset is already known at runtime, this information can be used to bias the simulated annealing. In experiments, the value of performance model predictions were mostly in the interval $[-7,7]$. Known invalids are set to $-10^{6}$, removing them from the top-$n$ ranking. Known valids are handled with more care. Since we do not know their performance, a place in the top-$n$ cannot be guaranteed. Therefore, we bias known valids with a value of $+1$. This boosts their chance to place in the top-$n$ ranking without overshadowing potentially better implementations.

All experiments in this section use pre-recorded ground truth data from the VTA hardware, the same as in the search space analysis evaluation in Section III. They act as proxys to real hardware measurements. This reduces the time to evaluate individual candidates and enables repeated experiments for robustness evaluation. Details on the operators in the working data set are in Table II.

Before the actual tuning process begins, the presampling algorithm from section IV-A explores the search space and gathers the data to replace random initialization of the first measure epoch $E_{0}$ with a set of configurations created by the process from section IV-B. From there, the existing flow of AutoTVM is used: measurements in epoch $E_{i}$ are used to train the performance model. This performance model is then used by the simulated annealing algorithm to select the configuration for next measure epoch $E_{i+1}$. This new batch of configurations is then evaluated on hardware and the cycles repeats. The original AutoTVM implementation has been altered in two locations: 1) The simulated annealing now uses the bias information from section IV-C and 2) hardware evaluations are replaced with querying the ground truth data gathered once for the full search space. All experiments use an epoch size of 50 measurements, or trials, until a total of 750 evaluations on hardware are performed. In all experiments with the extended auto-tuning, the presampling set $N\subset S$ of each search space $S$ is built form $min(1000,|S|)$ configurations. Model initialization is then performed as described in Section IV: $E_{0}=E^{valid}\cup E^{invalid}$, with $E^{valid}=min(|N^{valid}|,25)$ and $E^{invalid}=min(|N^{invalid}|,50-|N^{valid}|)$.

First, the effects of the improved model are investigated. Runs with randomly initialized search and an unmodified AutoTVM serve as a baseline. The results for the individual runs are in Figure 5. The improved methods outperform the baseline for every workload, with a mean of only $0.416\times$ as many trials for convergence. Two results stand out: layer 59 is remarkably better ($0.111\times$) than the mean and layer 107 shows weaker improvements ($0.735\times$). These are also the workloads with lowest (layer 59) and the highest (layer 107) numbers of valids in the search space. For layer 59, if there are very few valids, finding any of these has much higher chance of being the best one. Layer 107 on the other hand, has the highest number of valids candidates, meaning that finding valids alone is not sufficient for good performance, making the search process harder. Nonetheless, our method improved the convergence time by $1.36\times$. Across all experiments, the mean IQR is only $0.424\times$ of the baseline. In some layers, like 3, 5, 8, 42, 48 or 92, IRQ measure was improved by factors of three to six.

All the median values of the presented method are consistently below 300 trials, with only layers 8, 17 and 53 requiring more than 200 trials. Further, the robustness of the proposed method roughly follows the baseline, but always improves upon it. Meaning that a layer with narrow IQR in the baseline, will also have a narrow IQR with the improved initialization and a wide IQR in the baseline leads to a wider IQR in the new method. While the overall performance improved, some runs still produced outlier performing significantly worse than the rest, for example layers 8, 42 or 53. Interesting is layer 42, where the IQR with presampling is narrow, but several positive and negative outliers exist. The baseline, on the other hand spreads from 50 to 750 trials, with IQR between 200 and 450. This hints towards a workload, that is rather difficult for the performance model to reason about and where the initial dataset has a large impact on the rest of the search. This is the complete opposite of layers 3, 5 or 106, where the spread is narrowed considerably and outliers are few or non-existent. The high robustness of 59 can again be explained with extremely low number of valid configuration in the search space. Often the best candidate was in the initial training batch.

In the previous experiments, the AutoTVM model used random sampling to start the search process and used an uninitialized model, whereas the presampling method already provided validity information for the start of the tuning process. This showed to significantly search performance and robustness. In the following experiments AutoTVM with a performance prediction model trained on 30000 random samples from all workloads in Table II is compared against the presampling based tuning. Since the model only learned about other workloads, this feature is called transfer learning.

The experiments with transfer learning will use the context relation features, as they provide the ability to learn across different operators [6]. From the fully evaluated search spaces, 30000 configurations were randomly sampled as training data (but never from the workload under evaluation). From these, ca. 63.04% could be used to train the model. The other samples proved to be invalid configurations during feature extraction phase and thus were removed from the training set. Although the context relation features are designed to produce feature vectors compatible across workloads, the vectors of some workloads still required zero padding to fit the input length. Due to the large overheads during transfer learning, experiments were only repeated 10 times, instead of 20.

Figure 6 compares the convergence time of this work and the trained tuner. Overall, the pretrained model is only needs 90.5% of the trials to find the best result, compared to the AutoTVM baseline. However, in almost half of the experiments, the tuning performance decreased. Five workloads (3, 5, 78, 92, 106) perform worse than the random baseline, layer 5 needs $2.085\times$ as many trials in the median. On the other hand, layers 8 or 42 performed significantly better than the presampling based method by factors of $2.24\times$ and $1.54\times$. For the robustness, transfer learning improves the baseline by a factor of $1.107\times$, but has strong variations in both directions. Layers 8 and 42, achieve an exceptional robustness, with the IQR being only 2.7% and 6.5% of the baseline. For layer 59, learning from other workloads is not successful, as IQR increases by a factor of $5.386\times$, although the convergence is slightly better.

Overall, the pretrained model consistently requires over 200 samples for median convergence, with the exceptions of layers 42, 53 barely and 107. Layers 8, 42, 53 and 107 show high robustness and no outliers, while outperforming the pre-samping method. However, workloads 5, 78, 92 or 106 are the complete opposite of this. They need over 300, the latter three over 500 samples to converge to the best solution. They also have a low robustness with the outer quartiles reaching up to the sampling limit. While these results improve upon the random baseline, our presampling based method outperforms the transfer learning approach in more than half the cases and never drops below baseline performance.

The previous sections compared the tuning performance by the number of trials, not time-to-solution. This metric makes comparison of different approaches easier and decouples the evaluation from runtime effects. However, these methods also introduce overheads in the overall process. While hardware execution is a major contributor to the overall tuning time, other overheads can increase the absolute runtime. Table III lists the per run overheads of the different methods. Performing the validation of individual candidates for model initialization has a significant overhead, while the sample selection overhead is the negligible. For layers 78, 8 and 92 presampling plus sample selection is still 82.6s, 415.4s and 246s faster than transfer learning. However, the presampling has to be performed for every new workload, whereas a once trained model can be reused for different workloads. Tuning the knob features, which just flatten the configuration into a vector, is faster than the curve features by factors of $8.44\times$, $14.42\times$ and $12.89\times$ for the layers in Table III. The data shows that for actual tuning runs, both the pre-tuning and runtime overhead are lower with our method, compared to the transfer-learning solution of TVM. Even when using an already trained model, the presampling overhead is amortized against the much more expensive feature extraction.