[CVPR 2021] UP-DETR: Unsupervised Pre-training for Object Detection with Transformers

안녕하세요, 허재연입니다. 오늘 리뷰할 논문은 CVPR 2021 및 TPAMI 2022에 게재된 UP-DETR이라는 논문으로, DETR에 대한 unsupervised learning을 다룬 논문입니다. 리뷰 시작하겠습니다.

Introduction

Self-Supervised Learning, Unsupervised pre-training은 사람이 annotation한 GT 라벨 없이 raw data만을 가지고 데이터에 대한 표현력을 확보하고자 합니다. 보통 2단계로 나뉘는데, 첫 번째 단계에서는 pretext task를 이용하여 대규모 데이터셋으로 모델을 사전학습 시키고, 두 번째 단계에서는 적용하고자 하는 downstream task에 사전학습한 모델을 fine-tuning 하게 됩니다. 연구자들은 좋은 pretext task(contrastive learning 등)를 고안하여 pretext task에 대한 성능이 아닌, 이후 적용시키고자 하는 downstream task를 잘 수행할 수 있도록 모델에 좋은 representation을 학습시키는 것을 목표로 합니다. 이러한 representation learning은 최근 자연어 처리 분야와 컴퓨터비전 분야 모두에서 학습 리소스 및 annotation에 대한 부담을 완화시킬 수 있을 것이라는 기대를 받으며 활발히 연구되었고, (일부 downstream task에서는 supervised learning 방법으로 사전학습한 성능을 능가하는)의미 있는 발전을 이루었습니다.

Transformer가 등장한 이후, 자연어처리 분야 뿐만 아니라 컴퓨터 비전에서도 attention 및 encoder-decoder 구조를 활용하려는 시도가 계속되었고, ViT나 Swin Transformer, DETR 등과 같은 모델이 개발되었습니다. 이 중 DETR(Object detection with transformers)은 object detection task를 트랜스포머 인코더-디코더를 통한 direct prediction problem 관점으로 바라보고 개발된 object detector입니다. 기존의 object detection에는 물체의 localization과 classification을 둘다 잘 수행해야했기에 복잡한 내부 구조를 가지고 있었고 이에 따라 사람이 직접 튜닝해주어야 하는 요소(anchor box, IoU threshold, model confidence 등)가 많았습니다. 하지만 DETR은 이러한 handcrafted 요소들(proposal, sample selection, non-max suppression과 같은 hand-designed 요소들) 없이 기존 detector의 성능에 견줄 수 있게 만들어졌습니다. DETR에 대해서는 이전에 다른 연구원분들이 자세히 리뷰해주신적이 있으니, 자세한 구조가 궁금하신 분들은 참고하면 좋을 것 같습니다: [이상인 연구원님 리뷰] [조현석 연구원님 리뷰]. 이런 DETR에도 transformer 구조를 채택함으로써 나타나는 단점이 있었는데, 그 중 하나는 대규모 학습 데이터 및 매우 긴 학습 시간이 필요하다는 것입니다.

하나의 예시로 COCO와 같은 데이터셋은 그나마 학습 데이터가 많긴 하지만, 이보다 학습 데이터가 적은 PASCAL VOC 데이터셋에서는 DETR이 충분히 잘 동작하지 못하는 경향이 있다고 합니다. COCO 데이터셋에 대해서도 DETR이 수렴하는데에는 Faster RCNN보다 10배 이상의 학습 주기가 필요하다고 알려져 있습니다.

좋은 비지도 사전학습 방법이 개발됨에 따라 사전학습 모델은 자연어처리(GPT, BERT등) 뿐만 아니라 컴퓨터비전(MoCo, SwAV 등)에서도 좋은 결과를 보여주었습니다. DETR은 CNN backbone(일반적으로 ResNet50. ∼23.2M parameters)을 사용해서 feature map을 추출한 다음 transformer encoder-decoder를 겨쳐 마지막에 prediction head로 예측을 하는 과정을 거칩니다. 이 때 CNN backbone은 사전학습된 가중치를 사용하여 좋은 visual representation을 추출할 수 있지만 transformer module(~18.0M parameters)은 사전학습되지 않았었다고 합니다. 많은 visual representation learning(contrastive learning 등) 방법들이 최근 관심을 받고 연구가 진행되었지만 (당시에) 현존하는 pretext task들로는 DETR의 transformer를 곧바로 학습시킬 수는 없었기 때문입니다.

저자들은 NLP에서 unsupervised learning의 큰 성공(BERT)에 영감을 얻어, DETR의 트랜스포머를(현존하는 pretext task로는 이를 사전학습시키기에 적절하기 않기에) 대규모 데이터셋(e.g ImageNet)에서 unsupervised방법으로 사전학습시켜 object detection task를 downstream으로 활용하고자 합니다. 이에 random query patch detection 이라는 이름의 새로운 unsupervised pretext task를 활용해 Unsupervised Pre-training DETR (UP-DETR)라는 비지도 사전학습 기법을 제안합니다. 해당 방법론은 주어진 이미지에서 무작위로 여러 query patch들을 crop하고, detection을 위한 트랜스포머가 주어진 이미지에서 이러한 쿼리 패치들의 bounding box를 예측하도록 사전학습합니다. 이러한 사전학습 과정을 통해 저자들은 2가지 중요한 이슈를 다루었다고 합니다 :

1. Multi-task learning : 모두 알고 계시듯 object detection은 객체에 대한 classification 및 localization task의 결합입니다. query patch detection이 classification을 위한 feature를 손상시키는것을 방지하기 위해 저자들은 사전학습 backbone을 freeze하고, 트랜스포머의 feature discrimination을 보존하기 위해 patch feature reconstruction을 도입하였습니다.
2. Multi-query localization: 다른 object query들은 다른 position area와 박스 크기에 집중합니다. 이를 묘사하기 위해 우선 간단한 single-query 사전학습을 제안하고 이후 이것을 multi-query 버전으로 확장합니다. multi-query patch들에 대해서는 query patch들과 object query들 간의 assignment 문제를 해결하기 위해 object query shuffle 및 attention mask를 도입하였다고 합니다.

실험 결과, UP-DETR은 PASCAL VOC 및 COCO object detection에서 더 빠른 수렴과 더 나은 average precision을 보였습니다. 뿐만 아니라, UP-DETR은 panoptic segmentation 등에서도 좋은 성능을 잘 transfer한다고 합니다. Ablation 실험적 결과로는 사전학습 CNN-Backbone을 freeze하는게 사전학습 과정에서의 feature discrimination 보존에 중요하다는 것을 발견했다고 합니다.

Method

UP-DETR또한 기존의 unsupervised visual representation learning 기법들과 같이 2단계로 구성됩니다. (1) 트랜스포머를 GT가 없는 대규모 데이터셋(실험에서는 ImageNet)으로 비지도 사전학습 시키고, (2) 모델 전체를 labeled data를 가지고 downstream task로 fine-tuning을 진행합니다. random query patch detection을 활용해 어떻게 트랜스포머의 인코더와 디코더를 사전학습 시키는지에 집중해서 살펴보도록 하겠습니다.

Encoder

우선 주어진 input image에 대해 frozen CNN backbone이 visual representation을 CxHxW tensor feature map 형태로 추출합니다. 그리고 이렇게 추출된 feature map에 positional encoding(2D)를 더하고 DETR의 multi-layer transformer encoder에 전달됩니다.

Decoder

본래 DETR과 약간 다르게, 사전학습 과정에서 input image로부터 query로 사용할 랜덤 패치를 crop하고 이 쿼리의 좌표를 저장합니다(이 정답 정보를 가지고 사전학습을 진행합니다). 이 random cropped query patch는 Global Average Pooling이 있는 CNN backbone을 거쳐 path feature p가 되고, p는 flatten되어 DETR의 디코더에 전달되기 전에 object query q에 더해집니다(q와 p는 차원이 동일하겠죠. object query는 원래 DETR의 디코더의 입력으로 사용되는 요소이고, query path는 비지도학습을 위해 저자들이 새로 도입한 요소이니 헷갈리면 안됩니다). 이후 디코더는 입력 이미지에서 random query patch들의 위치를 예측하도록 훈련됩니다. 논문에서는 이해를 위해 간단히 Single-query patch 상황을 가정하여 먼저 설명하고, 이후 multi-query patch에 대해 설명합니다.

Single-query patch

DETR에서 각 object query들은 서로 다른 공간적 요소들을 학습하게 됩니다. 각각 서로 다른 object query들이 다른 위치, 지역 정보와 박스 크기 등에 집중하게 되는 것이죠(논문에서는 각각이 집중하는 부분에 대해 ‘different spatial specialization’이라고 표현합니다). 해당 방법론에서는 이미지에서 패치를 랜덤하게 crop하여 쿼리 패치에 대한 위치나 박스 크기 등의 사전 정보가 없습니다. 위 Figure2의 (a)에는 모든 object query들(N=3)에 대한 하나의 single-query patch(M=1)가 나타냈습니다.

사전학습 과정에서 patch feature p는 각각 다른 object query q에 더해지고, 디코더는 입력 이미지에서 query patch의 bounding box를 검출하기 위해 N pair의 predictions를 생성합니다. DETR의 학습 매커니즘과 동일하게 저자들도 예측 집합과 GT 값을 학습하는데 동일한 match cost 및 헝가리안 알고리즘을 사용하였습니다. 모든 match pairs에 대한 Hungarian loss는 다음과 같습니다 :

수식에서 에서 C_i는 쿼리 패치 매칭에 대한 이진 분류를 수행하며(ci가 0 또는 1), L_cls는 2개 클래스에 대한 cross-entropy입니다. b_i는 박스의 좌표 및 너비 높이{x,y,w,h}를 정의하는 벡터이며, L_box는 L1 loss와 generalized IoU(GIoU loss)를 선형 결합한 loss term입니다. L_rec는 사전학습 과정 도중 classification과 localization에 대한 밸런스를 유지하기 위해 본 논문에서 제안한 reconstruction loss입니다. 해당 loss는 바로 아래에서 살펴보겠습니다.

Patch Feature Reconstruction

object detection은 서로 다른 feature를 필요로 하는 object classification과 object localization이 결합된 task입니다. 저자들은 feature reconstruction term L_rec을 제안하여 localization 사전학습 과정에서 classification feature가 보존되도록 하였습니다. CNN에서 추출된 feature discrimination을 feature가 트랜스포머에 전달된 이후에도 보존되게 하고자 한 것입니다. L_rec는 CNN backbone에서 추출된 patch feature를 L2정규한 다음 mean squre error 취한 것으로, 다음과 같이 정의됩니다. p_i는 ResNet50 백본의 reconstruted feature(C=2048)입니다

Frozen Pre-training Backbone

저자들의 실험 결과 CNN backbone은 트랜스포머의 사전학습에 매우 크게 영향을 미쳤다고 합니다. CNN backbone을 트랜스포머와 함께 사전학습하여 이 둘을 처음부터(scratch) random query patch 검출로 함께 학습하게 하면 잘 수렴하지도 못하고 성능도 떨어졌다고 합니다. 해당 문제는 DETR에서도 동일하게 나타납니다.

또한, object detection은 서로 다른 feature 특성을 필요로 하는 object classification과 localization가 결합된 task인데, 저자들의 사전학습 전략은 feature classification 대신 localization만 고려하기에 사전학습 도중에는 CNN backbone을 freeze하였다고 합니다.

Multi-query patch

실제로 일반적인 object detection 상황에서는 각 이미지에 여러 object가 있습니다(COCO의 경우 장당 평균 7.7개의 object가 있다고 합니다). 이 뿐만 아니라 single-query path는 object query N이 커질 때 수렴에 어려움이 생길 수 있다고 합니다. 실제 검출 상황 및 사전학습 방법의 효율성을 고려했을 때, single-query patch만을 이용하는 것보다 multi-query patch를 활용하는 방향으로 확장하는 것이 자연스럽습니다. 저자들은 N개의 object query들을 M개 그룹으로 나누어 각 query path가 N/M object query들에 할당되도록 하였습니다(Fig.2(b) 참고). 이 때 각 query patch들은 완전히 랜덤하게 crop되어 서로 상관관계가 없도록 하였습니다.

Attention Mask

쿼리 패치 간 독립성을 위해 저자들은 attention mask matrix를 활용해 다른 object query들 간 상호작용을 조정했다고 합니다. N*N차원의 mask matrix X는 디코더의 self-attention의 softmax layer에 다음과 같이 더해집니다.

attention mask는 다음과 같이 정의됩니다 :

X_i,j 는 object query qi_i가 object query q_j과의 interaction에 참여하는지를 결정합니다. 직관적으로 표시하기 위해 Fig.2에는 0과 -∞을 각각 1과 0으로 표시하였습니다. 이 Attention Mask는 기존에 제안되었던 UniLM의 token mask와 비슷하다고 하는데, 제가 해당 방법론은 잘 모르기에 object query를 각 random query patch에 잘 할당하기 위해 사용하였다고 이해하였습니다.

X_i,j 는 object query qi_i가 object query q_j과의 interaction에 참여하는지를 결정합니다. 직관적으로 표시하기 위해 Fig.2에는 0과 -∞을 각각 1과 0으로 표시하였습니다. 이 Attention Mask는 기존에 제안되었던 UniLM의 token mask와 비슷하다고 하는데, 제가 해당 방법론은 잘 모르기에 object query를 각 random query patch에 잘 할당하기 위해 사용하였다고 이해하였습니다. 사전학습 과정에서 dropout과 비슷한 효과를 주기 위해 query patch들의 10%를 0으로 randomly masking하였다고 합니다. 저자들의 실험에서는 N=100, M=10으로 설정되었다고 합니다.

Experiment

저자들은 UP-DETR을 ImageNet에서 사전학습하고 VOC 및 COCO에서 object detection, one-shot detection, panoptic segmentation task에 대해 파라미터를 파인튜닝했습니다. 모든 실험에서 ResNet50 backbone을 사용하였고, 6개의 트랜스포머 인코더 및 디코더를 사용하였습니다. CNN backbone은 SwAV로 사전학습된 가중치를 활용하였고, 주어진 이미지에 대해 쿼리 패치를 랜덤 좌표로 crop해 너비와 높이를 128×128픽셀로 resize한 뒤 augmentation을 적용하였습니다(증강 기법은 horizontal flip을 제외하면 SimCLR와 동일하게 적용하였다고 합니다)

PASCAL VOC Object detection

완전한 비교를 위해서 DETR뿐만 아니라 Faster R-CNN(R50-C4 백본 사용)이 함께 리포팅되었습니다. 일단 DETR은 Faster RCNN과 비교해 PASCAL VOC에서 크게 성능이 떨어집니다(논문 서두에서는 DETR이 완전히 수렴할만큼 PASCAL 데이터셋이 크지 않기 때문이라고 합니다). UP-DETR은 이러한 DETR과 비교해 성능을 크게 끌어올렸으며, AP에 대해서는 Faster RCNN보다 좋은 성능을 기록했습니다. Fig3에서는 VOC 데이터셋에 대한 AP(COCO style) learning curve를 나타내었는데, UP-DETR이 매우 빠르게 수렴하는 것을 확인할 수 있습니다. learning rate를 줄인 이후에도 UP-DETR이 DETR보다 계속 높은 성능을 유지합니다. 300epoch보다 더 길게 학습을 하여도 DETR은 up-detr을 따라집지 못했습니다. 저자들은 이에 대해 트랜스포머를 사전학습하는것이 학습 데이터가 부족한 상황에서 필수불가결하다고 합니다.

COCO object detection

COCO 데이터셋에는 작은 물체가 많습니다. 그리고 DETR은 small object detection 성능이 부족하다는 문제를 가지고 있습니다(DETR의 고질적인 문제인 긴 수렴 시간과 작은 객체 탐지 성능이 떨어지는것을 보완하기 위해 이후 Deformable DETR이 제안되었죠). 이 때문에 저자들은 정확한 비교를 위해 AP, AP50, AP75, AP_s, AP_m, AP_L을 함께 나타내었습니다. 해당 표에도 포괄적인 비교를 위해 잘 최적화된 Faster-RCNN(FPN)을 나타내었습니다.

UP-DETR은 AP, AP_M, AP_L에서 가장 좋은 결과를 ㅂ여주었습니다. DETR에는 FPN과 같은 구조가 없기에 작은 물체에 대한 검출 성능이 떨어져 AP_S에 대해서는 Faster RCNN만큼 좋은 성능을 기록하지는 못했습니다. Fig.3b는 COCO 데이터셋에 대한 AP learning curve를 나타내었습니다. VOC만큼 큰 격차는 아니지만, 여전히 UP-DETR이 더욱 빠른 수렴을 보여줍니다.

Panoptic Segmentation

detection 뿐만 아니라 segmentation에 대한 실험도 수행되었습니다. COCO dataset에 대해서 수행되었으며, DETR보다 약간 좋은 성능을 보이거나 근접한 성능을 보이는, 좋은 결과를 달성하였습니다.

Ablataion

CNN을 freeze하는 것과 feature reconstruction을 활용하는 것에 대한 ablation이 수행되었습니다. CNN을 freeze하지 않았을 때는 up-detr의 성능이 저하되었으며, 심지어 DETR보다도 낮은 성능을 기록했습니다. feature reconstruction도 많은 성능 향상을 기록했기에 UP-DETR은 이 둘을 함께 사용하는 방향으로 설계되었습니다.

Conclusion

random crop을 활용한 unsupervised learning 방법이 처음은 아니었지만, 지난번에 읽었던 논문은 ResNet backbone 및 contrastive learning을 활용한 방법이었습니다. DETR의 특성에 맞게 pretext task를 잘 설계한 것 같네요. 논문을 읽으며 DETR에 대한 이해도가 약간 더 깊어진 것 같습니다.

감사합니다.