[ICLR 2024] CLIPSelf: Vision Transformer Distills Itself for Open-Vocabulary Dense Prediction

안녕하세요. 이번 주 X-Review에서는, 24년도 ICLR에 Spotlight으로 게재된 <CLIPSelf: Vision Transformer Distills Itself for Open-Vocabulary Dense Prediction>이라는 논문을 소개해드리겠습니다. ICLR Spotlight으로 지정된 논문인데, 통찰력 있는 문제 정의와, 이를 해결하는 메소드를 굉장히 간단히 구성하면서도 성능을 많이 끌어올려 좋은 평을 받게 되었다고 생각합니다.

1. Introduction

오늘 소개해드릴 논문은 Open-Vocabulary Dense Prediction task를 수행합니다. Dense Prediction은 영상의 픽셀 단위로 학습하거나 예측하는 task를 의미하는데, Object Detection (OD)이나 Semantic Segmentation (SS), Depth Estimation 등이 이에 속할 수 있습니다. 본 논문에서는 OD와 SS 성능을 통해 제안하는 방식의 효과를 보여주고 있습니다.

최근 Closed-set에서의 OD와 SS는 Real-world에 존재하는 무한한 visual concept에 대응할 수 없다는 단점때문에 Open-Vocabulary 상황에서의 연구가 활발히 진행되고 있다는 점을 모두 알고계실 것입니다. 초기의 Open-Vocabulary 관련 연구에서는 Image-text 레벨로 학습되어있는 CLIP의 지식을 통해 기존 task를 Open-Vocabulary로 확장하고 있었습니다. CLIP의 Global representation 관련 지식을 local representation에 주입시켜 Dense Prediction을 수행할 수 있게 되는 것입니다.

저자는 ViT 기반 CLIP 모델의 지역적 feature를 Dense feature라고 칭하는데, 이 Dense feature의 region-language alignment 능력을 평가하기 위해 위 그림 1과 같은 사전 실험을 하나 진행합니다. 그림 1-(a)는 ImageNet에서 ResNet 기반 CLIP과 ViT 기반 CLIP의 이미지 분류 성능을 보여주고 있습니다. 전반적으로 ResNet보다 ViT가 단일 영상에 대한 표현력이 더 좋음을 알 수 있네요.

이후 그림 1-(b)에선 COCO 데이터셋의 Region classification 성능을 보여주고 있습니다. 다들 아시겠지만 COCO는 ImageNet과 다르게 한 영상 내 여러 객체가 등장하죠. 그래서 이에 대해 Image Crop과 Dense Feature, 두 상황에서의 분류 정확도를 측정합니다. Image Crop은 영상 내 객체를 crop하여 해당 영역의 global feature를 분류에 사용하는 상황이고, Dense Feature는 영상의 각 패치가 표현되어있는 feature에서 객체 영역만을 따와 이를 분류에 활용한 상황입니다. ViT의 Dense Feature 분류 정확도가 크게 하락하는 것을 통해, ViT 기반 CLIP은 결국 좋은 Global representation을 가지고 있지만, 패치 레벨의 표현력은 굉장히 떨어진다고 해석할 수 있습니다.

마지막으로 그림 1-(c)는 Dense Feature를 시각화해본 정성적 결과로, 위와 같은 해석의 이유를 보여주고 있습니다. 시각화는 K-Means clustering을 통해 feature가 유사한 패치를 엮어 붉은색 또는 푸른색으로 표시한 것입니다. 가장 우선적으로 CLIP의 사전학습 방식 자체가 Global token과 텍스트에 대한 Contrastive learning이기 때문에 이러한 결과가 나왔을 수도 있습니다.

또한 시각화 결과를 보면, ResNet 기반 CLIP은 기본적으로 CNN 자체가 inductive bias를 갖추고 있기에 영상의 local한 영역을 잘 구분하여 물체와 배경에 대한 나름의 구별력을 갖추고 있습니다. 그러나 ViT 기반 CLIP 모델의 Dense Feature는 전체 영상 하나를 표현하는 데에 특화되어있고, Self-attention을 통해 모든 패치 토큰들이 서로 상호작용을 하기에 영상의 local한 표현력을 전혀 파악하고 있지 못합니다. 이런 Dense Feature로 Dense prediction을 수행하면 좋은 결과를 기대하기 어려울 것입니다.

결국 저자가 이야기하는 기존의 문제는, “ViT 기반의 CLIP은 영상 전체는 굉장히 잘 표현하지만, local한 영역을 보는 능력은 떨어지기 때문에 Dense Prediction을 수행하기엔 적합하지 않다”입니다. 이러한 문제를 해결하기 위해 저자는 CLIP 모델이 자기 스스로로부터 지식을 전수받는 간단한 학습을 하나 수행하는데, 이를 따 방법론의 이름을 CLIPSelf라고 칭합니다. 자세하게는 방법론 파트에서 말씀드리겠습니다.

2. Methodology

2.1 Image Representation vs Dense Representation

CLIPSelf 방법론 자체가 CLIP의 Image Representation(한 영상 전체를 표현하는 벡터)과 CLIP의 Dense Representation(토큰별 벡터로 구성되는 한 영상의 feature) 간 학습을 수행하기 때문에 먼저 각각이 무엇인지 정의하고 시작합니다.

CLIP’s Image Representation

모델의 마지막 attention block을 거쳐 나온 출력을 x = (x_{0}, x_{1}, \cdots{}, x_{h \times{} w}^{T}와 같이 표현할 수 있습니다. 여기서 x_{0}은 클래스 토큰을 의미하며 나머지는 h \times{} w개의 토큰별 벡터를 의미합니다. 이후 하나의 영상을 표현하는 Image Representation x_{image} = z[0]은 아래 수식에 따라 얻을 수 있습니다. Self attention을 거쳐 기존의 클래스 토큰을 한 번 업데이트하고 사용하는 것이죠.

CLIP’s Dense Representation

CLIP ViT로부터 h \times{} w 형태의 feature map \mathcal{X}_{dense}를 얻기 위해 아래와 같은 과정을 거칩니다.

위 수식에서 Emb()는 FC layer, Proj()는 Conv layer에 해당합니다. 이후 클래스 토큰에 해당하는 z'[0]을 버리고 z'[1:h \times{} w] 형태의 이미지 임베딩을 h \times{} w 형태로 reshape하여 \mathcal{X}_{dense}를 얻고, 이후 영역에 대한 feature를 따온 뒤 RoIAlign을 거쳐 형태를 맞춰주게 됩니다.

2.2 CLIPSelf

앞서 Image Representation과 Dense Representation의 추출 방식에 대해 간단히 알아보았습니다. 저희가 예상했던대로 각각은 feature의 클래스 토큰과 영역 feature에 해당했는데, 이제 어떠한 대상에 대해 위와 같은 feature를 추출하고, 어떤 학습을 진행하는지 살펴보겠습니다.

우선 CLIPSelf는 그 이름에 걸맞게 CLIP 모델이 자기 자신으로부터 self-distillation을 하는 방법론에 해당합니다. 앞서 Introduction에서 말씀드린 바에 따르면, CLIP 모델의 Image Representation은 표현력이 좋지만 Dense Representation은 상대적으로 부족하였습니다. 그림 2-(b)에서 CLIPSelf의 아이디어를 살펴볼 수 있습니다.

우선 그림 2-(a)의 기존 Dense Prediction 학습 방법론에 비해 CLIPSelf는 영상의 텍스트 캡션 쌍을 필요로 하지 않으며 RPN을 통해 추출한 foreground 영역만을 학습하는 기존 방법론에 비해 CLIPSelf는 영상의 모든 영역을 학습함으로써 다양한 배경에 대한 강인성도 기를 수 있게 됩니다. OD는 foreground에 해당하는 물체만을 잡으면 되니 배경에 대한 강인성을 기르는 것이 간접적인 장점이 될 수 있지만, Segmentation의 경우 ‘stuff’에 해당하는 뒷배경 또한 잡아야하니 직접적인 도움이 될 수 있습니다.

Image Patches as Regions.

Self-distillation을 위해 두 개의 기존 CLIP 모델을 가져옵니다. Teacher와 Student 모델 모두 CLIP의 사전학습 가중치로 초기화되며, Student만이 CLIPSelf의 방식을 통해 fine-tuning되게 됩니다. 이를 위해 영상을 우선 m \times{} n 형태의 패치로 쪼개줍니다. m과 n은 \{1, \cdots{}, M\} 중 하나로 임의 선택되며, 실제로는 M=6이라고 합니다.

Self-Distillation.

우선 Student 모델은 영상 하나를 통째로 입력받아 앞서 설명드린 방식대로 \mathcal{X}_{dense}를 추출합니다. 여기서의 \mathcal{X}_{dense}는 표현력이 상대적으로 떨어지는, 학습해야하는 대상이라고 볼 수 있습니다. 이렇게 추출한 dense feature에서 각 패치에 해당하는 구역을 따오고 Pooling하여 \mathcal{P}^{ij} (i \in{} \{0, \cdots{}, m-1\}, j \in{} \{0, \cdots{}, n-1\})을 얻을 수 있습니다. 위 그림에서 Dense Feature에 색으로 나눠져있는 6개 부분을 각각 풀링하여 Dense feature에서 각 패치의 feature를 추출하는 것입니다.

그리고 앞서 m \times{} n 형태로 잘라놓은 원본 영상의 각 패치는 Freeze되어있는 Teacher에 입력되어 각각의 Global Representation을 추출합니다. 단순하게 각 패치를 한 영상으로 보고 이에 대한 x_{image}를 추출하는 것과 동일한데, 여기서는 각 패치의 feature를 t_{image}^{ij}라고 칭합니다.

이후에는 두 feature의 코사인 유사도를 최대화하는 방향으로, 위 수식을 통해 학습을 진행합니다. 위 수식이 CLIPSelf가 제안하는 유일한 loss에 해당합니다. Dense Prediction에 반드시 필요하지만 표현력이 부족한 s_{dense}^{ij}이 표현력 좋은 t_{image}^{ij}로부터 distillation을 받는 것입니다.

3. Experiments

3.1 Application to Open-Vocabulary Dense Prediction

Experiment Setting.

CLIPSelf 방식으로 학습된 CLIP ViT Backbone은 Open-Vocabulary OD와 SS 수행에 활용됩니다. 기본적으로 타 방법론들과의 fair comparison을 위해 동일한 데이터셋 스플릿으로 학습을 진행하였다고 합니다. 예를들어 OVOD의 COCO면 COCO의 train2017 split, LVIS 데이터셋이면 train split, OVSS의 COCO면 COCO의 train2017 split과 같이 학계와 동일한 세팅으로 학습을 진행했다는 의미입니다.

OVOD

위 표 3에서 OVOD 성능을 확인할 수 있습니다. 기본적으로 CLIPSelf의 효과를 확인하기 위해 F-VLM이라는 기존 방법론의 scheme을 따라갑니다. F-VLM은 이미 학습되어있는 ResNet 기반 CLIP을 freeze해두고 뒷단의 detector를 학습한 방법론입니다. 이 때 ResNet 기반 CLIP을 ViT로 변경한 것을 F-ViT라 칭하고, 이 ViT를 CLIPSelf 방식으로 학습한 ViT로 교체한 것이 F-ViT+CLIPSelf에 해당합니다. 이에 대해선 기존 F-ViT보다 훨씬 높은 성능을 달성하고 있고, ResNet 기반 모델보다 더욱 높은 성능을 보여주고 있습니다.

OVSS / OVPS

표 4, 5는 각각 OVSS와 OVPS에서의 성능입니다. 앞선 OVOD와 동일하게 CLIPSelf는 백본의 표현력을 향상시키는 방법론이다보니 기존 SS와 PS를 수행하는 방법론의 백본만을 CLIPSelf ViT로 교체하고 측정한 성능에 해당합니다. SS와 PS 모두에서 기존보다 개선된 성능을 보여주고 있습니다. 데이터셋에 따라 편차는 있지만 특히 SS에서 큰 성능 향상 폭을 보여줍니다.

Comparison with Using Region-Text Pairs

위 표 6은 RegionCLIP이라는 기존 방법론과의 성능 비교 표입니다. Region-Text 학습을 위해 RegionCLIP이라는 방법론에서는 RPN을 통해 얻은 foreground 영역과 object noun을 매칭시켜 region-text pair를 만들고 학습에 사용하게 됩니다. 이러한 경우 우선 모델 구조 내 RPN을 필요로 하며, 해당 영역과 매칭시켜줄 대상인 텍스트를 필요로 한다는 추가적 cost가 발생하게 됩니다.

Fair comparison을 위해 CLIPSelf도 Region proposal을 수행하였으며 결과를 봤을 때 noisy한 region-text pair를 만들어내는 Region-Text 방법론보다 모든 task의 모든 평가지표에서 더욱 높은 성능을 보이는 것을 확인할 수 있습니다.

4. Conclusion

해당 논문을 통해 기존 CLIP이 어떠한 특성을 가지고있었는지 알아볼 수 있었으며, 이를 해결하기 위해 외부에서 지식을 가져오는 것이 아닌 CLIP 모델 자기 자신으로부터 지식을 전수받는다는 점이 인상깊은 논문이었습니다. 또한 추가적 cost 없이 굉장히 간단한 방식으로 백본의 표현력을 크게 향상시킬 수 있다는 포인트가 ICLR의 Spotlight이 될 수 있었던 이유라고 생각합니다.

이상으로 리뷰 마치겠습니다.