{'fraction_non_alphanumeric': 0.045436300875010036, 'fraction_numerical': 0.0201894517138958, 'mean_word_length': 4.949140401146132, 'pattern_counts': {'":': 0, '<': 0, '<?xml version=': 0, '>': 0, 'https://': 1, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
Searching for TrioNet: Combining Convolution with Local and Global Self-Attention

Recently, self-attention operators have shown superior performance as a stand-alone building block for vision models. However, existing self-attention models are often hand-designed, modified from CNNs, and obtained by stacking one operator only. A wider range of architecture space which combines different self-attention operators and convolution is rarely explored. In this paper, we explore this novel architecture space with weight-sharing Neural Architecture Search (NAS) algorithms. The result architecture is named TrioNet for combining convolution, local self-attention, and global (axial) self-attention operators. In order to effectively search in this huge architecture space, we propose Hierarchical Sampling for better training of the supernet. In addition, we propose a novel weight-sharing strategy, Multi-head Sharing, specifically for multi-head self-attention operators. Our searched TrioNet that combines self-attention and convolution outperforms all stand-alone models with fewer FLOPs on ImageNet classification where self-attention performs better than convolution. Furthermore, on various small datasets, we observe inferior performance for self-attention models, but our TrioNet is still able to match the best operator, convolution in this case. Our code is available at https://github.com/phj128/TrioNet.


Introduction
Convolution is one of the most commonly used operators for computer vision applications [21,22,49]. However, recent studies [14,26,45,58,71] suggest that the convolution operator is unnecessary for computer vision tasks, and that the self-attention module [56] might be a better alternative. For example, Axial-DeepLab [58] built stand-alone self-attention models for image classification and segmentation. Although these attention-based models show promising results to take the place of convolution-based models, these attention-based models are all human-designed [45,58], making it challenging to achieve optimal results on new dataset, tasks, or computation budgets.   Given a target dataset, task, and computation budget, Neural Architecture Search (NAS) [39,40,47,48,55,64,74,75] is an efficient technology to automatically find desirable architectures with marginal human labor. Due to its effectiveness, previous works [7,40,47,69,74] have successfully applied NAS to different computer vision tasks, including object detection [17,59], semantic segmentation [38], medical image analysis [73], and video understanding [62]. However, most of the searched models for these computer vision tasks are built upon convolutional neural network (CNN) backbones [24,25,51].
Previous convolution-based NAS methods usually consider attention modules as plugins [36]. During the search procedure, the search algorithm needs to decide whether an attention module should be appended after each convolution layer. This strategy results in the searched network architecture majorly consisting of convolution operators, causing a failure to discover the potentially stronger attention-based models. Therefore, in this paper, we study how to search for combined vision models which could be fully convolutional or fully self-attentional [45,58]. Specifically, we search for Tri-oNets where all three operators (local-attention [45], axial-attention [58], and convolution) are considered equally important and compete with each other. Therefore, self-attention is no longer an extra plugin but a primary operator in our search space. Seemingly including the self-attention module into the search space could achieve our goal, we observe it is difficult to apply the commonly used weight-sharing NAS methods [6,7,69] due to two issues.
One issue is that self-attention modules and convolution blocks make our search space much more complicated and imbalanced than convolution-only spaces [6,7,69]. For instance, convolution usually contains kernel size and width to search while self-attention have more options like query, key and value channels, spatial extent and multi-head numbers. The self-attention operators correspond to much larger search spaces than convolution, with the same network depth. This imbalance of self-attention and convolution's search space makes the training of supernets intractable. Therefore, we propose Hierarchical Sampling, which samples the operator first uniformly before sampling other architecture options. This sampling rule suits our setting better because it ensures an equal chance for each operator to be trained in the supernet, alleviating the bias of search space size.
The other issue is that the multi-head design of attention operators poses a new challenge for weight-sharing [7,69] NAS algorithms. Current weight-sharing strategy [7,53,69] always shares the first few channels of a full weight matrix to construct the weight for small models. However, in self-attention modules, the channels are split into multi-head groups to capture different dependencies [56]. The current weight-sharing strategy ignores the multihead structure in the weights and allocates the same channel to different heads depending on the sampled multi-head groups and channels, forcing the same channel to capture different types of dependencies at the same time. We hypothesize and verify by experiments that such weight-sharing is harmful to the training of supernets. Instead, we share our model weights only if they belong to the same head in multi-head self-attention. This dedicated strategy is named Multi-Head Sharing strategy.
We evaluate our TrioNet on ImageNet [34] dataset and various small datasets [16,33,43,44,63]. The TrioNet architectures found on ImageNet are shown in Fig. 1. We observe that TrioNets outperform stand-alone convolution [21], local-attention [45] and axial-attention [58] models with fewer FLOPs on ImageNet where self-attention performs better than convolution. On small datasets where self-attention models perform inferior to convolution, our TrioNet is still able to matche the best operator with fewer FLOPs on average.
To summarize, our contributions are four-fold: (1) We regard self-attention and convolution as equally important basic operators and propose a new search space that contains both stand-alone self-attention models and convolution models. (2) In order to train a supernet on our highly imbalanced search space, we adopt Hierarchical Sampling rules to balance the training of convolution and self-attention operators. (3) A Multi-Head Sharing strategy is specifically designed for sharing weights in multi-head self-attention operators. (4) Our searched TrioNet reduces computation costs and improves results on ImageNet classification, compared with hand-designed stand-alone networks. The same phenomenon is observed when TrioNets are searched on small datasets as well.
2 Related work Self-attention. Self-attention was firstly proposed for NLP tasks [1,56,66]. People then successfully apply self-attention module to many computer vision tasks [2,8,11,29,61,72]. More recently, it has been shown that self-attention can be used to replace convolution in vision models. Hu et al. [26] and Local self-attention [45] build the whole model with self-attention restricted to a local patch. SAN [71] explores a boarder self-attention formulation. Axial-DeepLab [58] extends local self-attention to global self-attention with axial self-attention blocks. Later, ViT [14] approaches image classification with the original NLP transformer architecture. Variants [20,41,65,70] of ViT propose a few hand-designed ways of applying local constraints to the global self-attention in ViT. In this paper, we automatically search for an architecture in a combined space that includes convolution, local self-attention [45], and axial global self-attention [58].
Neural Architecture Search. Neural Architecture Search (NAS) was proposed to automate architecture design [74]. Early NAS methods usually sample many architectures and train them from scratch for picking up a good architecture, leading to large computational overhead [7,39,48,50,55,69,74,75]. More recently, people develop an one-shot pipeline by training a single over-parameterized network and sampling architectures within it to avoid the expensive train-from-scratch procedure [4,5,6,13,19,27,40,47,67]. However, the search space of these methods is restricted in the MobileNet-like space [24,25,51]. There are also some attempts to search for self-attention vision models [35,36,60]. These works mainly focus on the single block design [60] and positions [36], or searches in a small space [35]. In this paper, self-attention and convolution are considered equally in our search space.

Method
In this section, we first define our operator-level search space that contains both convolution and self-attention. Next, we discuss our one-shot architecture-level search algorithm that trains a supernet. Finally, we present our proposed Hierarchical Sampling (HS) and Multi-Head Sharing (MHS) that helps training the supernet.

Operator-Level Search Space
Convolution is usually the default and the only operator for a NAS space [6,7,55,69,74]. In this paper, however, we introduce self-attention operators into our operator-level space and search for the optimal combination of convolution and self-attention operators [56]. Specifically, we include efficient self-attention operators that can be used as a standalone operator for a network. We use axial-attention [58] instead of fully connected 2D self-attention [14] as an instantiation of global self-attention for computational efficiency.
Local Self-Attention. Local self-attention [45] limits its receptive field to a local window. Given an input feature map x ∈ R h×w×d in with height h, width w, and channels d in , the output at position o = (i, j), y o ∈ R d out , is computed by pooling over the projected input as: The relative positional encodings r p−o ∈ R d q are also learnable vectors and the inner product q T o r p−o measures the compatibility from location p = (a, b) to location o = (i, j). In practice, this single-head attention in Equ. (1) is extended to multi-head attention to capture a mixture of affinities [56]. In particular, multi-head attention is computed by applying N single-head attentions in parallel on x o (with different learnable matrices W n Q ,W n K ,W n V , ∀n ∈ {1, 2, . . . , N} for the n-th head), and then obtaining the final output z o by concatenating the results from each head, i.e., z o = concat n (y n o ). The choice of local window size m significantly affects model performance and computational cost. Besides, it is not clear how many local self-attention layers should be used for each stage or how to select a window size for each individual layer. For these reasons, we include local self-attention in our search space to find a good configuration. Axial Self-Attention. Axial-attention [23,29,58] captures global relations by factorizing a 2D self-attention into two consecutive 1D self-attention operators, one on the height-axis, followed by one on the width axis Both of the axial-attention layers adopt the multi-head attention mechanism, as described above. Note that we do not use the PS-attention [58] or BN [30] layers for fair comparison with local-attention [45] and faster implementation.
Despite capturing global contexts, axial-attention is less effective to model local relations if two local pixels do not belong to the same axis. In this case, combining axial-attention with convolution or local-attention is plausible. So we study the combination in this paper by including axial-attention into our operator-level search space.

Architecture-Level Search Space
Similar to hand-designed self-attention models, Local-attention [45] and Axial-DeepLab [58], we employ a ResNet-like [21] model family to construct our architecture-level search space. Specifically, we replace all 3×3 convolutions by our operator-level search space that contains convolution, local-attention, and axial-attention.  Figure 2: Supernet candidate sampling. Left: ProxylessNAS-style [6]. Middle: OFA-style [7]. Right: our TrioNet with shared projection layers and three separate parallel operators. Tab. 1 summarizes our search space for each block. Different from a cell-level search space in the NAS literature [7,69], our spatial operators are not shared in each residual block. As a result, our search space allows a flexible combination of different operators at each layer. This space also includes pure convolutional [21], pure local-attention [45], and pure axial-attention [58] ResNets. Our search space contains roughly 7.2 × 10 25 models in total, 3.6 M times larger than that of OFA [7] (2 × 10 19 ).

Searching Pipeline
Given this search space, we employ the one-shot NAS pipeline [7,69], where the entire search space is built as a weight-sharing supernet. The typical strategy [6,53] of most oneshot NAS works is shown in Fig. 2 a). A MobileNet-like [24,25,51] supernet is built with blocks that contain several parallel candidates with different channels and receptive fields. Consequently, the supernet needs a lot of parameters to hold the large search space, making it inapplicable to our huge search space (Tab. 1). OFA and BigNAS [7,69] propose to share candidate weights in a single block, as shown in Fig. 2 b). This allows a much larger search space with manageable parameters but it still limited one operator only. In our case, different spatial operators can not share all the parameters.
Based on these works [6,7,69], we make a step forward. We insert local- [45] and axial- [58] attention into the single block parallel with the spatial convolution as the primary spatial operator. These parallel spatial operators share the same projection layers and increase the flexibility of the supernet without introducing many parameters. Like [7,69], we can formulate the problem as min where (C (W o , α)) denotes that the selection of the parameters from the weights of the supernet W o with architecture configuration α, Γ(A) is the sampling distribution of α ∈ A and L train denotes the loss on the training dataset. After training the supernet, an evolutionary algorithm is used to obtain the optimal model α o on the evolutionary search validation dataset with the goal as Finally, we retrain the searched model from scratch on the whole training set and validate it on the validation set. Next, we discuss two issues of the default supernet training pipeline and our solutions.

Hierarchical Sampling
Previous work [57]  the spatial operator uniformly. Given the operator, we then randomly sample a candidate from the operator space. We can formulate this as α = (β , θ ), β ∼ U(B), θ ∼ U(Θ) (4) where the architecture α can be expressed as operators configuration β ∈ B sampled from Tab. 1 (the first row) and weights configuration θ ∈ Θ sampled from Tab. 1 (from the second to the sixth row) with uniform distribution U. In this way, we ensure an equal sampling chance for all operators.
In addition, we notice that the training is biased towards middle-sized models. We attributes this to the fact that the search space is full of middle-sized models and thus random sampling trains mostly middle-sized models. To address this issue, we use Sandwich rule [68,69] that samples the smallest candidate, the biggest candidate and 2 random candidates. Sandwich rule is adopted in our Hierarchical Sampling after the operator for each block is selected.

Multi-Head Sharing
The weight-sharing strategy for convolution has been well studied. Previous works [3,7,53,69] share common parts of a convolution weight for different kernel sizes and channels. However, this strategy cannot be simply applied to our search space with multi-head selfattention operators. The varying number of heads makes sharing channels more complicated than convolution. In multi-head attention, the channels are split into multi-head groups to capture different dependencies [56]. Thus, as shown in Fig. 3, the default weight-sharing strategy allocates the same channels for different multi-head groups, which is harmful to the supernet training. To deal with this issue, we propose Multi-Head Sharing (MHS). As shown in Fig. 3, we take into account the multi-head structure and first split all output channels into number-of-head groups. Then, we share channel weights only if they belong to the same head in multi-head self-attention. We show the selection procedure with PyTorch-like [46] pseudo code in Alg. 1. In this way, we ensure that different channel groups do not interfere with each other. Algorithm 1 Pseudo code of Multi-Head Sharing in a PyTorch-like style # w: weight matrix, c_in: in channels # c_out: max out channels, s_c_out: sampled out channels, n: multi-head number h_c = s_c_out / n w = w.reshape(n, -1, c_in) w = w[:, :h_c, :] # h_c, n, c_in w = w.reshape(s_c_out, c_in) return w

Experiments
Our main experiments are conducted on ImageNet [34] dataset. We also provide detailed ablation studies on the proposed Hierarchical Sampling and Multi-Head Sharing. Finally, we

ImageNet Classification
Tab. 2 shows the main results of our searched models with different FLOPs constraints. The searched models are compared with stand-alone convolution [21], stand-alone local selfattention [45] and stand-alone axial-attention [58] models. As shown in Fig. 4, with 2B FLOPs budget, our TrioNet-A achieves 74.8% accuracy and outperforms ResNet with 57.4% less computation. With 3B FLOPs budget, our TrioNet-B achieves 75.9% accuracy, which outperforms stand-alone local self-attention and axial-attention respectively with 33.3% and 14.3% fewer FLOPs. These results show that our searched TrioNet is able to outperform all hand-designed single-operator networks in terms of computation efficiency. In addition to our main focus in the low computation regime, we also evaluate models in the high computation regime. With 4B and 4.7B FLOPs budgets, our TrioNet-C and TrioNet-D achieve comparable performance-FLOPs trade-offs with stand-alone axial-attention and still outperforms fully convolution [21] or local self-attention [45] methods with fewer FLOPs. Compared with OFA [7] under the same training settings as ours, TrioNet outperforms OFA with 1.1%, 0.2%, 0.2% and 0.3% accuracy with 16.7%, 11.8%, 13.0% and 2.1% less computation. Note that the larger models (TrioNet-C and TrioNet-D) are already close to the limit of our architecture space, which might lead to performance degrade. If large models instead of lightweight models are desired, our TrioNet searching pipeline can be directly extended to a high computation search space as well.

Ablation Studies
In this subsection, we provide more insights by ablating each of our proposed components separately. The experiments are also performed on ImageNet [34]. We monitor the training process of the supernets by directly evaluating the largest possible model for each operator. Hierarchical Sampling. Fig. 5 a) visualizes the supernet training curves with and without our proposed Hierarchical Sampling. We observe that sampling all candidates uniformly hurts convolution models a lot due to the low probability of convolution being sampled. However, with our proposed Hierarchical Sampling, all the model results are improved. It is worth mentioning that with our Hierarchical Sampling strategy, the local self-attention [45] models are improved too, even if they are not sampled as often as the case without Hierarchical Sampling. We hypothesize that the comparable performances of all three parallel candidates help the optimization of all operators. Multi-Head Sharing. Fig. 5 b) compares the training process of the supernet with and without our proposed Multi-Head Sharing. We observe that our Multi-Head Sharing strategy helps large multi-head self-attention models achieve higher accuracy. Similar curves are observed for small models as well, though the curves are omitted in the figure. In addition to the training curves of the supernet, we also analyze the searched model results. As shown in Tab. 3, adopting Multi-Head Sharing in the supernet training improves the searched architecture by 1.5%, from 75.4% to 76.9%. Sandwich rule. Fig. 6 a) plots the training curves with and without Sandwich rule [68,69]. It shows that Sandwich rule helps the training of large models by a large margin for all operators. Fig. 6 b) shows the distribution of the sub-models selected from the supernet under different strategies. It can be observed that some middle-sized models are comparable with the biggest models because they are trained more and gain a lot from the large models optimization [69]. However, this accuracy partial order cannot reflect the training from scratch accuracy, and the searched model accuracy in Tab. 3 also shows this. With Sandwich rule, the sampling distribution is changed and the ranking of these models are kept. Number of epochs. Weight-sharing NAS methods require a long supernet training schedule because the candidates need to be well-trained before they can accurately reflect how good each candidate is. Similarly, we evaluate how our TrioNet searching algorithm scales with longer training schedules. As shown in Fig. 7 a), our searched models does not perform well with 20 epochs and 60 epochs, probably because the candidates have not been well-trained with such a short schedule. However, we do not observe a huge difference between 180 epochs and 540 epochs, probably because our supernet saturates with our simple ResNetlike training recipe and the weak data augmentation used. Amount of data. Our TrioNet searching algorithm is also tested with different amount of data. To achieve this goal, we train the supernets with 10k, 100k images and the full training dataset (we still remove the evolutionary search set). Then, the searched architectures are still trained from scratch on the full training set. In this way, we evaluate only the contribution of data on the searching algorithm, or the quality of the searched architecture, which is decoupled with the amount of data used to the searched models from scratch. Fig. 7 b) shows that the searching on only 10k images gives a poor architecture, and our searching algorithm scales well, i.e., finds better architectures, with more data consumed. Summary. Tab. 3 summarizes the searched model accuracies of our ablated settings. We notice that without Sandwich rules [68,69]

Results on Small Datasets
The goal of NAS is to automate the architecture design on the target data, task, and computation budget. From this perspective, we adapt our TrioNet NAS algorithm to other datasets beyond ImageNet [34], in order to evaluate if the algorithm is able to find a good architecture on the target datasets. These adaptation experiments are performed on five small datasets: Stanford Cars [33], FGVC Aircraft [43], CUB [63], Caltech-101 [16] and 102 Flowers [44]. Overall Performance. Tab. 4 shows the results of our TrioNet models searched directly on various small datasets [16,33,43,44,63] and finetuned from supernet weights, as well as the accuracies of baseline models. Empirically, we notice that axial-attention [58], which outperforms ResNet [21] by a large margin on ImageNet [34], performs poorly on these small datasets with a big gap to ResNet (8.6% accuracy on average), probably because the global axial-attention uses less induction bias than convolution or local-attention. However, in this challenging case, our TrioNet finds a better architecture than the hand-designed local self-attention [45] and axial self-attention [58]

Results on Segmentation Tasks
In this section, we evaluate TrioNet on semantic segmentation [9,42] and panoptic segmentation [32] tasks. The ImageNet pretrained models are employed. For semantic segmentation, we apply DeepLabV3 [10] on PASCAL VOC datasets [15]. For panoptic segmentation, we perform the experiments on COCO datasets [37] with Panoptic-DeepLab [12] under a short training schedule. All experiments only replace the backbone with TrioNet-B.
Results. Tab. 6 shows the semantic segmentation and panoptic segmentation results. Using TrioNet-B as the backbone outperforms Axial-26 [58]

Conclusion
In this paper, we design an algorithm to automatically discover optimal deep neural network architectures in a space that includes fully self-attention models [45,58]. We found it is not trivial to extend the conventional NAS strategy [7,69] directly because of the difference between convolution and self-attention operators. We therefore specifically redesign the searching algorithm to make it effective to search for self-attention vision models. Despite our observation is from studying the self-attention module, we believe this is readily to extend to searching for other components such as normalization modules [30].