[ICLR 2022] Open-Vocabulary Object Detection via Vision and Language Knowledge Distillation

오늘 가져온 논문은 Open-Vocabulary Object Detection 분야의 논문입니다. 일반적인 Detection 모델과는 달리, 임의로 주어지는 text input 에 해당하는 object를 이미지 내에서 찾는 task 이죠.
현존하는 Object Detection Dataset 으로는 class의 수를 scale-up 하기에 어렵기 때문에 본 논문에서는 CLIP과 같은 사전학습된 open-vocabulary image classification model 을 teacher로 삼은 뒤 two-stage detector 모델을 student 로 설정하여 학습을 진행합니다.

1. Introduction

본 논문에서 설계하는 Open-Vocabulary Object Detection (OVOD) 모델의 컨셉은 아래 Figure 1에서 잘 보여주고 있습니다. 크게 Image Dataset에 라벨링 되어있는 Base Category와, text 정보만 가지고 있는 Novel Category 로 나누어지며, Base Category 만으로 학습된 모델일지라도, 분류하고자 하는 Novel Category 의 text 정보만 있다면 아래처럼 검출이 가능합니다.

사실 검출하고자 하는 클래스의 수를 늘리는 직관적인 해결책은 Detection Dataset에 포함된 클래스의 수를 직접적으로 늘리는 것입니다. 실제로 LVIS Dataset은 1,203개의 클래스를 가지고 있죠. 하지만 이는 long-tail distribution, 즉 클래스 불균형 문제가 불가피합니다.

반면, Detection Dataset 말고 Image-Text pair dataset의 경우는 인터넷 상에 매우 풍부하게 분포되어 있습니다. 대표적으로 CLIP 도 이러한 dataset을 사용해서 vision-language model 을 구축하였죠. 이런 CLIP과 같은 pretrained open-vocabulary classification model의 image-level representation 능력을 빌려서, 본 논문의 저자는 OVOD 를 수행하고자 합니다.

저자들은 우선 OVOD task를 크게 2가지 sub-optimal problem으로 구분짓습니다.

1. General Object Proposal
2. Open-Vocabulary Image Classification

우선 Base Category 정보를 가지는 Image Dataset을 사용해서 Region Proposal Network를 학습시킵니다. 여기서 중요한 점은, 해당 RPN 모델은 Base Category뿐만 아니라 Novel Category에 대한 proposal 도 잘 예측해 낼 수 있어야 합니다. 이후 검출된 Base+Novel proposal에 대해 CLIP과 같은 Open-Vocabulary Image Classification Model 을 적용해서 전체 category를 예측할 수 있습니다.

저자들은 LVIS 데이터셋에서 드물게 등장하는 category를 novel category로 설정하고, 나머지를 base category 로 설정합니다. 놀랍게도 novel category에 대한 성능이 supervised 성능을 능가했다고 말합니다만, 검출된 object proposals 영역을 하나하나 classification model 에 입력으로 넣어서 class 를 예측해야 하기 때문에 추론 속도가 매우 느리다는 단점이 존재합니다.

이러한 문제점들을 해결하고자 저자는 ViLD (Vision and Language knowledge Distillation) 모델을 제안합니다. 이는 크게 2가지 요소로 구성되어 있습니다.

1. ViLD-text: learning with text embeddings
2. ViLD-image: learning with image embedding

우선 ViLD-text 에서는 사전 학습된 text encoder에 category name 을 넣어서 text embedding을 얻게 됩니다. 이때 base category만을 사용하여 학습을 진행합니다.

ViLD-image 에서는 detector 모델이 검출한 object proposal 영역을 사전 학습된 image encoder에 넣어서 image embedding (teacher)을 얻게 됩니다. 그리고 detector 모델로부터 추출되는, proposal 영역에 대한 region embedding (student)이 있을텐데, 이를 image embedding 과 align 시키는 distillation 과정을 진행하게 됩니다. ViLD-text와는 달리 base+novel category 에 대한 학습을 진행하게 됩니다.

---

2. Method

Notations

OVOD 실험을 위해 detection dataset을 base category와 novel category 로 나눠야 합니다. 나눠진 subset을 각각 C_B, C_N 으로 표기합니다. 학습 시에는 C_B 의 annotation 만 사용하게 되겠죠. 또한 CLIP 과 같은 사전학습된 open-vocabulary image classifier 의 text encoder와 image encoder를 각각 \mathcal{T, V} 로 표기합니다.

2.1. Localization for Novel Categories

OVOD 수행에 있어서 첫번째 문제 사항은 novel objects의 위치를 잘 찾아야 한다는 것입니다. 이를 위해 Mask R-CNN 과 같은 기존 2-stage detector 모델의 구조를 조금 변형했다고 합니다. 일반적으로 이러한 모델 내에서 localization을 담당하는, bounding box regression과 mask prediction layer 는 class-specific한 방식으로 동작합니다. novel category 에 대한 localization 수행을 위해 해당 두 요소를 class-agnostic하게 변형하였다고 합니다. category별로 하나씩이 아닌, 모든 category에 대해 단 하나의 bbox와 mask를 예측하도록 말이죠. 이를 통해 novel objects 에 대해 일반화할 수 있다고 합니다.

2.2. Open-Vocabulary Detection with Cropped Regions

앞선 2.1 절 단계를 통해 object candidates 영역의 위치를 찾은 후, 사전학습된 open-vocabulary image classifier 를 사용해서 각 영역에 대한 분류를 수행해야 합니다.

Image embeddings

proposal network 를 base category C_B를 통해 학습시킨 후, region proposals \tilde{r}\in\tilde{P} 를 offline 으로 생성합니다. 그리고 해당 영역을 전체 이미지 I에서 crop & resize 한 후, 사전학습된 image encoder \mathcal{V}에 넣어서 image embeddings \mathcal{V}(crop(I,\tilde{r})) 를 계산합니다. 이때 이미지를 1\times, 1.5\times 크기로 각각 crop하여 추출된 image embedding을 앙상블 하여 사용한다고 합니다. 1.5\times의 경우 더 많은 context cues 를 제공한다고 하여, 앙상블 후 unit norm 으로 정규화 시킵니다.

Text embeddings

각 category text를 “a photo of {category} in the scene” 라는 prompt template을 사용하여 \mathcal{T} 에 입력으로 넣어 text embedding을 offline 으로 추출합니다. 여러 prompt 나 동의어가 제공될 경우 이를 앙상블하여 사용했다고 합니다.
이후, 생성된 image embedding과 text embedding 사이의 cosine similarity 를 계산하고 softmax 연산과 class별 NMS 과정을 거쳐서 최종 detection 결과를 얻습니다. 이때 intro에서 말씀드린것 처럼, crop된 영역에 대해 하나하나 모두 \mathcal{V} 에 입력으로 넣어서 class를 예측하기 때문에 추론 속도가 매우 느리다고 합니다.

2.3. ViLD: Vision and Language Knowledge Distillation

앞선 느린 추론 속도를 해결하기 위해 본 논문에서는 ViLD 라고 하는 기법을 제안합니다.
우선 base category로 학습된 two-stage detector을 사용해서 proposal r을 예측한 후, \mathcal{R}(\phi(I),r) 을 통해 region embedding 을 생성합니다. \mathcal{R}, \phi는 각각 backbone과 region embedding을 생성하는 lightweight head 를 의미합니다. classification layer 이전의 output 을 사용하여 region embedding 으로 사용한다고 합니다.

Replacing classifier with text embeddings

ViLD 의 구성요소 중 첫번째, ViLD-text 에 대한 설명입니다. 저자들은 사전 추출된 text embedding을 활용하여 분류를 수행하기 위해 region embedding을 학습시키는 것입니다. 위 Fig 3의 (a)->(b) 의 변경점을 살펴보시면 좋을 듯 한데, 기존 모델의 구조에서 마지막 classifier layer를 제거하고, text embedding 으로의 투영을 진행하는 projection layer를 부착한 후 projection된 region embedding 과 미리 추출된 text embedding 사이의 유사도 계산을 기반으로 한 분류가 수행됩니다.

학습 단계에서는 base class로 구성된 subset C_B 만을 사용하여 text embedding \mathcal{T}(C_B) 를 추출하여 region embedding 과의 유사도를 계산하게 됩니다. 이때 base category의 경우 \mathcal{T}(C_B) 와의 유사도 계산을 진행하게 되고, base category 이외의 region proposal에 대해서는 background 로 구분되도록 학습을 진행한다고 합니다. 이때, ‘background’라고 하는 단어 자체가 실제 background proposal 정보를 모두 잘 담고 있는 것이 아니기 때문에, 학습 가능한 embedding e_{bg} 를 사용한다고 합니다.

각 proposal region r 에서 계산된 region embedding e_r과 category별 text embedding t_i사이의 cosine similarity 를 계산하고 softmax 를 거친 후 cross-entropy loss 를 계산하게 됩니다. 아래 수식과 같습니다.

N은 image별 proposal의 수를 의미합니다.

학습 과정에서는 base category에 대한 text embedding 만 사용하지만, 예측 과정에서는 novel category를 포함하여 추출된 text embedding \mathcal{T(C_B\cup{C_N})} 을 사용한다고 합니다.

Distilling image embeddings

ViLD의 두번째 구성 요소, ViLD-image 에 대한 설명입니다. 이는 image encoder \mathcal{V}를 teacher로 삼아서 해당 지식을 student detector에게 distillation 시키는 과정입니다. 즉, student에서 추출된 region embedding \mathcal{R}(\phi(I),\tilde{r}) 을 teacher로 부터 추출된 image embedding \mathcal{V}(crop(I,\tilde{r}) 에 정렬(align)시키는 것입니다.

학습의 효율성을 위해, 각 학습 이미지에 대해 M개의 proposals \tilde{r}\in\tilde{P} 를 offline 으로 추출한 후, 이에 대한 M개의 image embedding을 미리 계산해둡니다. 앞서 말씀드렸다시피, region proposal network는 class-agnostic하게 학습되기 때문에, 추출한 M개의 proposal은 C_B, C_N 을 모두 포함하게 됩니다. 그리고 아래의 L1 loss를 통해 region embedding과 image embedding 사이의 distillation 이 진행됩니다.

그리고 앞선 두 loss를 더한 최종 ViLD 의 loss function은 아래와 같습니다. 또한 전체 모델의 구조 컨셉은 Fig.3-(d) 를 보시면 됩니다.

ViLD-image 에서의 distillation 과정은 학습 단계에서만 계산되며, 학습&평가 단계에서의 모델 동작 방식은 Fig.2를 참고하시면 됩니다.

2.4. Model Ensembling

본 논문에서는 base + novel category 모두에 대해서 최고의 detection 성능을 달성하기 위해 모델 앙상블 기법을 추가적으로 설계합니다. 우선 ViLD-text detector의 예측과 open-vocabulary image classfication model의 예측을 결합합니다. 이에 대한 직관은 ViLD-image 에서 distillation loss 를 통해 teacher model의 예측 (즉 open-vocabulary image classfication model 의 image encoder 로 부터 추출된 image embedding) 과 유사해지도록 학습이 진행되는데, 그렇다면 그냥 teacher (open-vocabulary image classfication model) 를 그대로 사용하는 것이 성능적으로는 더 좋겠지요.

ViLD-text detector를 사용해서 top k 개의 candidate region과 그에 해당하는 confidence score를 구해냅니다. category i에 대한 proposal \tilde{r}의 confidence score 를 p_{i, ViLD-text} 라고 표기합니다. 그리고 해당 region 영역을 전체 image I 에서 crop 하고 open-vocabulary classification model 에 입력으로 넣어서 구한 teacher 모델의 confidence score 를 p_{i, cls} 라고 표기합니다. 이 두 score를 각 base & novel category 에 대해 서로 다른 가중치를 적용하여 최종 앙상블 score 를 구해냅니다. 아래 과정을 통해 말이죠.

이때의 가중치 \lambda 는 2/3 를 사용해서 base category 에 대해서는 ViLD-text 를 더 많이 반영하고, novel category에 대해서는 open-vocabulary classification model 을 더 많이 반영하고자 하였습니다.

---

3. Experiments

Benchmark Settings

주로 LVIS v1 데이터셋을 사용하여 벤치마크를 진행했다고 합니다. LVIS 데이터셋은 OVD 에 적합한 데이터셋으로, 전체 1,203 개의 category 로 구성되어 있다고 합니다. 이중 자주 등장하는 866개를 base category C_B로 지정하고, 드물게 등장하는 337개를 novel category C_N으로 지정했다고 합니다. rare(novel) category 에서의 AP 인 AP_r 을 main metric으로 사용했다고 하네요.

추가적으로 COCO-2017 데이터셋에서는 WordNet hierarchy 이슈(?)로 인해 15개 category를 제외하고, 48개의 base category와 17개의 novel category로 실험을 진행했다고 합니다.

Learning Generalizable Object Proposals

base category만으로 학습된 detector가 novel category 에 대해서도 localize를 잘 하는지를 평가하는 실험입니다. ResNet-50 백본의 Mask R-CNN 의 RPN network 에 대해 평가를 진행하였고, novel category 에 대한 average recall (AR) 을 측정하였습니다.

base + novel category 모두를 보고 학습한 성능 대비, base category 만을 보고 학습한 모델의 recall 이 약 2정도밖에 하락하지 않은 것을 확인할 수 있습니다. 물론 unseen category generalization 에 초점을 두고 모델을 개선할 수 있었지만, 이는 추후 연구로 남겨두었다고 하네요.

Open-Vocabulary Classifier on Cropped Regions

표 2는 detector 에 의해 검출된 region proposal을 기존 이미지 I에 crop 해서 open-vocabulary classifier를 통해 분류한 Average Precision (AP) 성능을 리포팅 한 결과입니다. novel category 에 대한 성능인 AP_r 을 보시면 될 것 같습니다. 또한 open-vocabulary classifier로는 CLIP 을 사용했다고 합니다.

proposal 의 유무 판단을 CLIP의 confidence score 만으로 진행했을 때 13.0의 AP 를 달성합니다. 저자들은 CLIP의 confidence score가 localization의 질과는 관련성이 없다는 것을 실험적으로 발견하고, proposal 의 유무 판단을 위해 CLIP의 confidence score와 detector의 proposal objectness score를 기하평균 내었고, 결과적으로 18.9 의 성능을 달성했다고 합니다.

놀라운 점은 아래쪽 base+novel을 모두 사용하여 학습한 supervised, 그리고 category의 balance를 고려한 supervised-RFS 보다도 AP_r 이 꽤나 큰 차이로 높은 것을 볼 수 있습니다. 이를 통해 CLIP과 같은 open-vocabulary classifier 모델이 novel object를 분류하는데에 있어서 좋은 teacher 가 될 수 있음을 보입니다.

Vision and Language Knowledge Distillation

Text embeddings as classifiers (ViLD-text)

CLIP을 통해 추출된 text embedding 과 backbone->lightweight head를 거쳐 생성 region embedding 사이의 코사인 유사도를 기반으로 분류를 수행하는 ViLD-text 에 대한 실험입니다. text-only corpus 만으로 학습된 GloVe text embeddings 를 통한 AP_r 성능이 3.0인데에 반해, image-text pair 로 학습된 CLIP 을 사용했을 때 10.1의 AP_r 를 보이고 있네요.

Distilling image embeddings (ViLD-image)

다음은 ViLD-image 에 대한 평가입니다. cropped region proposals 영역에 대한 CLIP image embedding 으로 부터 distillation 을 통해 학습된 부분이죠. object category label 없이 학습을 했음에도 불구하고 11.2의 AP_r 성능을 보입니다. 하지만 나머지 AP 성능들은 제일 낮은 결과를 보이고 있네요. 이를 통해 visual distillation 과정이 open-vocabulary detection 에는 잘 동작하였지만, CLIP on cropped regions 에 비해서는 낮은 성능을 보인다는것을 알 수 있습니다.

Text+visual embeddings (ViLD)

앞선 ViLD-text 와 ViLD-image 를 결합한 결과입니다. 16.1의 AP_r의 성능으로, novel category 에 대한 성능이 상당 수준 향상하였습니다.

Model ensembling

2.4 절에서 언급했던 앙상블의 결과입니다. 앙상블은 크게 2가지 방식으로 진행됩니다.
1) ViLD-text with CLIP
2) ViLD-ensemble

이 중 첫번째는 cropped region에 대해 CLIP image encoder 를 하나하나 태우는 방식입니다. 그래서 당연히 성능적으로는 가장 높습니다만, 추론 속도가 느리기 때문에 real world에서의 실시간성 측면에 위배됩니다.

두번째는 ViLD-ensemble, 즉 2.4절에서 언급한 내용입니다. 성능 향상도 유의미하고, 속도 측면도 해결 가능합니다.

Stronger teacher model

Image Encoder 부분, 즉 teacher 모델을 더 무겁고 좋은 것으로 변경했을 때의 성능 향상입니다. 기존 실험에서 사용했던 CLIP ViT-B/32 대비 더 좋은 CLIP ViT-L/14와 ALIGN 을 사용하였다고 합니다. distillation 과정에서의 성능 향상이 꽤나 유의미하게 눈에 띄며, 이때 teacher 모델만 변경했기 때문에 추론 속도에는 변함이 없다고 합니다.

Qualitative Results

정성적인 결과입니다. 전체 4개의 행 중 2번째 행의 base+novel을 제외하고는 모두 학습 시에 보지 못한 category 에 대한 예측을 수행 한 결과입니다. 꽤나 잘 찾는것을 확인할 수 있습니다.

On-the-fly interactive object detection

단순 category text 를 입력으로 넣는것이 아닌, 더 자세한 fine-grained category 혹은 attributes(특징) 과 같은 임의의 text를 입력으로 넣었을 때의 결과를 보여주는 실험입니다. 본 저자들이 설계한 ViLD가 CLIP을 통해 단순 category 정보를 넘어서 fine-grained한, 특징적인 정보까지 학습했음을 보이는 결과입니다. (항상 좋은 결과를 보이는 것은 아니라고 하네요)

---

CLIP이 적용된 논문을 처음부터 끝까지 제대로 정독해본 건 이번이 거의 처음인 거 같습니다. 앞으로 text쪽 분야도 배척하지 않고 열심히 열린 마음으로 읽어야겠습니다.

감사합니다.