[Paper Review] Momentum Contrast for Unsupervised Visual Representation Learning

aaaa1313 2023. 3. 9. 23:42

1. Introduction

CV에서는 Unsupervsied Representation Learning에 대한 연구가 부족하였다. Unsupervised Representation Learning이 상당히 발달된 NLP의 경우, 토큰화된 사전을 구축할 때 별개의 단어, 문자와 같은 신호 공간을 사용하기 때문에 사용하기에 용이하지만, CV의 경우 data가 연속적이면서 high-dimensional 공간에 있기 때문에 Unsupervised Representation Learning을 사용하기가 어려웠다.

Unsupervised Visual Representation Learning을 위해 몇몇 논문들에서는 contrastive lsos를 사용하였다. 이를 사용한 목적들은 다 제각각이지만, 공통적으로 dynamic dictionary를 사용하였다. image나 patch의 data들을 encoder를 통해 dictionary 안에 있는 key라는 token과 query를 생성하였고, key와 query의 similarity를 이용하여 contrastive loss를 최소화 하고자 하였다.

저자는 이러한 점에서 dictionary가 $(i)$ large $(i i)$ consistent 한 것이 이상적이라 생각하고 있다. 왜냐하면 연속적이면서 high-dimensional한 CV에서는 dictionary가 크면서, key가 유사하거나 동일한 encoder를 통해 query와의 비교가 일관적인게 좋다고 생각하기 때문에 이러한 방식을 활용하는 Momenturm Contrast, MoCo를 제안하였다.

MoCo는 queue라는 dictionary를 통해 data sample들을 가지고 있으며, 앞서 언급했던 것과 같이 large하면서 consistent하다. 또한, 학습 과정에서 mini-batch로 나누어 학습할 때, 가장 오래된 것을 버리면서 가장 최신 mini-batch를 받아들이는 과정을 반복하면서 consistent하게 진행되도록 한다.

2. Related Work

2.1 Loss functions

MoCo에서 사용되는 contrastive loss는 similarity를 확인한다. 많은 연구에서 사용되는 loss의 경우, train된 model의 prediction과 고정된 target 간의 비교를 통해 이루어지지만, contrastive loss의 경우 target이 training 중 바뀔 수 있으며 network에 의해 계산된 data representation으로 정의될 수 있다.

2.2 Pretext task

일부 손상된 input을 recovering 해주거나 pseudo-labels를 형성하는 등 다양한 task가 존재한다.

3. Method

3.1 Contrastive Learning as Dictionary Look-up

encoding 된 query q와 dictionary 안에 있는 encoding 된 key들의 집합 $k_{0}, k_{1}, k_{2}, . . .$ 을 생각해보자. key들은 positive key( $k_{+}$ ) 한 개와 나머지 negative keys인 $k_{-}$ 로 구성되어 있으며, query와 유사한 $k +$ 는 contrastive loss가 작고, $k -$ 는 커진다.

contrastive loss는 queries와 keys를 represent하는 encoder networks를 training하기 위해 unsupervised objective function의 역할을 한다.

식 1에서 q와 k는 각각 $q = f_{q} (x^{q})$ , $k = f_{k} (x^{k})$ 로 표현 가능하며, $f (f_{q} = f_{k})$ 는 encoder를 나타내며, $x^{q}$ 와 $x^{k}$ 는 image나 patch, context 등을 나타낸다.

3.2 Momentum Contrast

MoCo에서는 good features를 배우기 위해서는 dictionary가 large하면서 consistent가 되어야 한다고 주장하였다. 이러한 방식을 활용하기 위해서 MoCo는 Mometum Contrast를 제안하였다.

3.2.1 Dictionary as a queue

Moco는 dictionary를 queue로 사용하면서 크기를 유지시킨다. mini-batch에서 가장 오래된 encoding된 keys를 제거하고, 제거한 만큼 새로 encoding 된 keys를 추가시킨다. 이러한 방법으로 크기를 유지시키며, queue의 크기는 mini-batch size와 관계가 없이 매우 크게 가지고 진행할 수 있다.

3.2.2 Momentum update

queue를 사용한다는 것은 dictionary를 매우 크게 사용할 수 있지만, key encoder에 대해 back propagation 하는것에 대해 어려움을 가질 수 있다. 왜냐하면 그만큼 sample수가 많고 sample에 대한 무수히 많은 gradient를 back propagation 을 진행해야 하기 때문이다.

처음에 제시한 해결책은 $f_{k}$ 에 대한 gradient 대신 $f_{q}$ 에 대한 gradient를 그대로 가지고 오는 것이다. 하지만 이러한 방법은 좋지 못한 결과를 야기시켰고, 이를 대체하여 식 2를 제안하였다.

$θ_{k}$ 와 $θ_{q}$ 는 각각 $f_{k}$ 와 $f_{q}$ 의 parameter이며, $θ_{k}$ 를 update 하는 방향으로 진행된다. m은 momentum coefficient이며, $θ_{q}$ 만 back propagation이 진행된다.

3.2.3 Relations to previous mechanisms

Fig 2는 contrastive loss 사용하는 mechanisms을 보여주고 있다.

a는 end-to-end 방식으로 현재의 mini-batch를 dictionary로 사용하고 있는 방법이다. 이러한 방법은 mini-batch 크기에 따라 dictionary 크기에 영향을 주기 때문에 mini-batch 크기가 매우 커야 한다. 이러한 방식은 GPU meomory size에 대한 문제가 생길뿐만 아니라, 특별한 network design을 요구하는 pretext task에도 적합하지 않다.

b는 memory bank를 사용하는 방식으로, memory bank란 dataset의 모든 samples의 representation을 저장하는 용도로 사용된다. 이러한 방식 또한 매우 큰 dictionary size가 요구된다. 또한, memory bank의 representation은 $f_{q}$ 에 의해서만 update가 되는데, key들은 다른 $f_{k}$ 에 의해서만 만들어 지므로 consistency가 떨어진다.

3.3 pretext task

MoCo에서 같은 image에서 나온 query와 key를 positive pair로 보고 다른 image에서 나온 quer와 key를 negative pari로 본다. 여기서 positive pair는 같은 image를 사용하지만 2개의 random view를 통한 augemntation을 진행한 data를 활용한다.

4. Experiments

4.1 Linear Classification Protocol

Fig 2에서 제안된 mechanisms 의 ImageNet linear classification을 보여주고 있다. end-to-end의 경우 GPU 제한으로 인하여 dictionary size를 더 크게 할 수 없어 결과가 제한됨을 확인할 수 있다. memory bank의 경우 end-to-end와 다르게 dictionary size를 크게 가져갈 수 있지만, MoCo보다 accuracy가 낮은것을 볼 수 있다.

위 Table은 ResNet-50을 mometum m의 값에 따른 accuracy를 확인할 수 있다.

Unsupervised learning methods에 대한 model과 model size가 각각 다르기 때문에, parameter를 기반으로 accuracy를 비교한 결과를 볼 수 있다.

4.2 Transferring Features

Unsupervised learning의 주된 목적은 transferrable 즉 obejct dection나 semantic segmantation과 같은 downstream task에서 적합하게 사용하기 위해서이다. 아래 tables에서는 앞서 언급한 tasks와의 비교를 확인할 수 있다.

Table 2는 PASCAL VOC에서 fine-tune된 Object detection 결과를 확인할 수 있으며, Table 3에서는 Fig 2에서 언급한 3개의 mechanisms에 대한 contrastive loss를 비교한 결과를 확인할 수 있다.

Table 5는 COCO로 fine-tune된 Obejct detection과 instance segmentation에 대한 결과를 확인할 수 있으며, Table 6에서는 이외의 다양한 donwstream tasks에 대한 결과를 확인할 수 있다.

5. Discussion and Conclusion

MoCo가 Unsupervised learning에서 좋은 결과를 보여줬지만, IN-1M이나 IG-1B와 같은 대규모 dataset에서 성능은 올라가지만 개선이 크게 되지 않는다는 것에 대한 discussion이 필요하다고 한다. 또한, MoCo가 instance discrimination task를 넘어서, masked auto-encoding과 language 등의 pretext task에서도 채택이 가능할 것으로 보고 있다.