본문 바로가기

Deep Learning/NAS

[2021.06] HR-NAS: Searching Efficient High-Resolution Neural Architectures With Lightweight Transformers

728x90

Image 단위의 Prediction을 하는 Classification Task와는 달리, Pixel 단위로 더 Finer한 Prediction을 하는 문제를 Dense Prediction이라고 합니다. Dense Prediction의 대표적인 예로는 Image Segmentation과 Pose Estimation 등이 있습니다. 본 논문은 Dense Prediction 문제를 푸는 Neural Architecture Search (NAS) 접근이 많이 없다는 문제 의식에서 시작을 합니다. Image Segmentation과 같은 Dense Prediction Task들은 Fine Detail의 정보를 담을 수 있는 High-Resolution (HR) Representation과 함께 Local Feature의 Ambiguity를 보상해 줄 수 있는 Global Context를 필요로 하는데요, 본 논문에서는 Global Context를 보는데 Multi-scale Convolutional Neural Network (CNN)만으로는 한계가 있기에, Transformer를 같이 사용하는 방법을 제안하고 있습니다. 여기서, Transformer의 연산 량을 줄이기 위해 Dynamic Down Projection이라는, CNN과 Plug-and-play 방식으로 결합될 수 있는 방법을 추가로 제안하고 있습니다. 여기에 더해 Multi-Scale CNN과 Transformer의 NAS를 용이하게 하기위해 다양한 실험을 통해 얻어진 Feature Normalization 방법을 소개합니다.

본 논문에서는 1) Dense Prediction과 같이 기존에 NAS가 많이 적용되는 않던 분야로 확장시킨 점과 2) CNN과 Transformer를 같은 Search Space안에 포함시켰다는 점 그리고 3) 경량 화된 Transformer의 하나의 방법정도를 제시하였다는 점 정도가 Key Contribution이라고 할 수 있을 것 같습니다.

Conference: CVPR 2021 (Oral)
URL: https://arxiv.org/abs/2106.06560
Code: https://github.com/dingmyu/HR-NAS

1. Methods

Our method aims to search for network structures within a multi-branch search space containing both Convolutions and Transformers with a resource-aware search strategy. In this section, we first introduce our lightweight Transformers. We then detail our multi-branch search space and how to integrate our Transformers into it. Finally, we describe the resource-aware fine-grained search strategy.

1.1. Lightweight Transformers

Figure 2. The lightweight Transformer, contains a projector, an encoder, and a decoder. It can be used plug-and-play to enhance the global context of image features.

기존에 Transformer는 Computation Resource가 크고, Resource 제약이 있는 상황에서는 사용할 수 없어서, 본 논문에서는 Lightweight Transformer를 제시하고 있다. (사실, Lightweight Transformer는 요즘 다른 논문들에서도 쉽게 찾아볼 수 있는 Topic 인 것 같다).  여기서는 Feature Dimension을 줄이기 위해 앞에 Projector가 들어가 있다는 점이 가장 큰 차이점이고, 더 간단하고 효과적인 2D positional mapping 방식을 제안하였다고 한다. 개인적으로, 이러한 Transformer 변형에 대한 제안이 어떠한 Novelty가 있는지는 잘 모르겠다.

1.2. Multi-Branch Search Space

"Inspired by HRNet [76], we design a multi-branch search space that contains both multiscale features and global contexts while maintaining high-resolution representations throughout the network."

Figure 3. (a) Our multi-branch search space, which is composed of parallel modules and fusion modules alternately. “1/4, 1/8, . . .” denote the down-sampling ratios. (b) Illustration of parallel modules and fusion modules. The red, black, and blue arrows represent the reduction searching blocks, normal searching blocks, and normal searching blocks with upsampling, respectively. The cubes represent feature maps. In this example, the fusion module generates an extra branch from the previous lowest-resolution branch by a reduction searching block. (c) The proposed searching block contains both convolutions with different kernel sizes C1, C2, C3, and a lightweight Transformer T.

SuperNet은 Searching Block이 Stacking 된 Parallel Module과 Multi-branch간 Information Exchange를 위한 Fusion Module로 구성되어 진다. SuperNet의 맨 앞에는 2개의 Convolution을 포함한 Input Conv Layers가 있어서, Feature Resolution을 1/4로 줄인다. 1/4로 줄어든 High-Resolution Branch에 더해서 차후에 Fusion Module에서 1/8, 1/16, 1/32의 Lower Resolution Branch가 추가되게 된다. 마지막 Layer에서, Multi-branch Feature들은 Resize후에 Concat하여 Final Classification Layer로 들어가게 된다. 

Fusion Block: For each output branch, all its neighboring input branches are fused by using the searching block to unify their feature map sizes. For example, a 1/8 output branch integrates information of 1/4, 1/8, and 1/16 input branches. In our fusion module, the high-to-low resolution feature transformation is realized by the reduction searching block, while the low-to-high resolution feature transformation is implemented with the normal searching block and upsampling.

Searching Block: As shown in Fig. 3 (c), the searching block contains two paths: one path is a MixConv [73], the other path is a lightweight Transformer for global contexts. The number of convolutional channels (rc) and the number of tokens (n) in the Transformer are searchable parameters. 마지막에 CNN과 Transformer에서 오는 두개의 Feature가 더해지는데, 이것은 기본적으로 CNN Path가 있고, Transformer Path에서 나오는 Global Attention Context가 Residual 하게 더해지는 것으로 볼 수 있다. 결국, 이 논문에서의 Novelty는 Transformer Path를 Plug-and-Play 방식으로 더해서, Global Context를 Encoding하는 General한 Searching Block을 제안했다 정도로 보면 될 것 같다.

1.3. Resource-aware Fine-grained Search

본 논문에서 Searching Block 내부에서 CNN의 Channel 갯수 (rc)와 Transformer의 Token (n)의 갯수를 Fine-grained하게 찾아가는 과정은, Network Slimming, AtomNAS와 DARTS에서 제안된 방법을 사용한다. Importance Factor α > 0를 Network Weight와 같이 학습해 나가면서, 학습 중간중간 마다 Importance Factor가 Threshold보다 낮은 Search Unit들을 제거해 나가는 방법이다. 다른 방법들과 같이, 제거 후에는 Batch-Normalization Layer의 Running Statistics를 Re-calibrate 해준다. 

Computation Cost가 높은 Unit들은 Importance Factor를 더 낮출 수 있도록, 아래와 같이 α에 Resource-Aware L1 penalty를 FLOPS 관점에서 제안하고 있다.

$\Delta_{i}=\left\{\begin{array}{ll}3 \times 3 \times h \times w, & i \in[0, r c) \\ 5 \times 5 \times h \times w, & i \in[r c, 2 r c) \\ 7 \times 7 \times h \times w, & i \in[2 r c, 3 r c) \\ O_{\mathcal{T}}\left(n^{\prime}\right)-O_{\mathcal{T}}\left(n^{\prime}-1\right), & i \in[3 r c, 3 r c+n)\end{array}\right.$

마지막으로 위에서의 Resource-aware Penalty Term을 추가하여, Loss Function을 제안하였다. 직관적으로는, 성능의 손실이 최소화되는 선에서 Computation Overhead가 높은 Search Unit들을 최대한 제거 (i.e., Importance Factor를 0으로) 하는 방향으로 설계되었다고 보면 될 것 같다. 

$L=L_{\text {task }}+\lambda \sum_{i \in[0,3 r c+n)} \Delta_{i}\left|\alpha_{i}\right|$

* Atomnas: Fine-grained end-to-end neural architecture search. In ICLR, 2020.
* Learning efficient convolutional networks through network slimming. In ICCV, 2017.
* DARTS: differentiable architecture search. In ICLR, 2019.

2. Experiments

Image Classification에 대한 실험만 놓고 봤을 때는, 다른 SOTA 방법들 대비 큰 장점이 부각되지 않는다. SCARLET-B와 ST-NAS-A는 ECCV 2020에 나온 논문으로, 참고할 가치가 있을 것 같다.

Figure 1. Comparisons on classification (ImageNet), segmentation (CityScapes), pose estimation (COCO), and 3D detection (KITTI), between the proposed approach and existing SoTA methods. Each method is represented by a circle, whose size represents the number of parameters. 

이외에도 논문에는 Semantic Segmentation (CityScape, ADE20K), Human Pose Estimation (COCO keypoint), 3D Object Detection (KITTI)에 대한 실험이 있다. 아무래도, Dense Prediction Task에 대해서는 NAS Benchmark들이 별로 없으니, 당연히 Global Context정보까지 같이 Search할 수 있는 본 논문이 비교우위가 있을 수 있을 것 같다.