[CVPR 2020] Momentum Contrast for Unsupervised Visual Representation Learning

[Link] Momentum Contrast for Unsupervised Visual Representation Learning

Self-Supervised 의 대표적인 방법론인 Moco를 리뷰하겠습니다.

---

Introduction

MoCo는 contrastive loss를 사용하는 self-supervised model 입니다.

당시 contrastive learning 을 기반으로 representation learning이 좋은 성능을 내고 있었는데요. 저자는 이 contrastive learning을 기반으로한 많은 방법론들이 결국 dictionary look-up 이라는 점에 주목하였습니다. 즉, 매우 많은 key가 존재하는 dictionary 중 query와 매칭되는 positive key는 높은 유사도를 가지고, negative key는 유사도가 낮아야 한다는 점에서 말이죠.

이 개념을 확장하여, 저자는 contrastive learning을 위한 dynamic dictionary를 구축하는 MoCo 를 제안합니다.

x_query 대한 representation q와 dictionary에 존재하는 representation k_0, k_1, k_2 ...에 대하여 같은 데이터로부터 발생된 키에 대하여 유사도를 높이고 다른 데이터로부터 파생된 키에 대하여 유사도를 낮추는 방향으로 학습을 진행합니다.

Method

본격적으로 MoCo의 방법론에 대하여 설명드리겠습니다.

Our hypothesis is that good features can be learned by a large dictionary that covers a rich set of negative samples, while the encoder for the dictionary keys is kept as consistent as possible despite its evolution.

본 논문에서 good features 를 위해 필요한 dictionary의 조건과 이유
=> “Large Dictionary” & “Consistent update Dictionary”

Contrastive Learning as Dictionary Look-up

Encoded 된 query q와 encoded 된 키k_0, k_1, ...이 있고 하나의 키 k_+만 q에 매칭되는 positive key 이고 나머지 K개의 키는 매치되지 않는 negative key 일 때, Equation 1으로 contrastive loss를 정의할 수 있습니다.

τ는 temperature 하이퍼 파라미터이며, 수식을 보면 고차원의 연속된 이미지에 대해 discrete한 dictionary를 구축하여 query에 대한 positive 샘플의 분류 확률을 높이도록 훈련된다고 해석할 수도 있습니다.

Momentum Contrast

MoCo에서는 dictionary를 크게 구성하기 위해 Queue로 구현하였습니다. 그 덕분에 mini-batch의 샘플을 negative 샘플로 지속적으로 사용할 수도 있고, 가장 오래된 샘플을 dictionary에서 제거하여 효율적인 딕셔너리의 key 관리가 가능하게 되었습니다. 따라서 dictionary의 샘플들은 천천히 바뀔 수 있었으며, 가장 오래된 샘플은 오래전 큐에 추가된 key representation으로 제거함을써 representation의 일관성을 유지할 수 있었습니다.

또한 Queue로 설계함으로써 딕셔너리의 사이즈를 배치 사이즈와 별개로 설정할 수 있었습니다. (이 점은 SimCLR 과 같은 논문과 비교하여 훌륭한 개선점이기도 합니다. 기존 논문들은 미니 배치 사이즈에 의존적이었고 더 키울 수 없었다는 한계가 있었습니다. 자세한 설명은 다음에 리뷰할 SimCLR 논문에서 설명드리도록 하겠습니다.) 아무튼 그 덕분에 dictionary의 크기를 설정할 수 있었고, 더 큰 key를 담을 수 있는 dictionary를 만들 수 있었다고 합니다.

그러나 학습을 위해서는 gradient back-propagation을 통해 key 인코더를 업데이트하는 것은 어려운 일이었습니다. (왜냐하면 큐에 담긴 모든 샘플에 그래디언트가 전달되는 것이 어려웠기 때문이죠). 이를 해결하기 위한 가장 간단한 방법은 f_q를 f_k로 복사하는 것이었는데, 이 방법은 dictionary의 일관성이 떨어져 성능이 좋지는 않았습니다. 따라서 MoCo에서는 Equation 2와 같이 momentum moving average로 f_k를 점진적으로 업데이트하는 방식을 제안합니다.

처음에는 f_k의 파라미터 θ_k와 f_q의 파라미터 θ_q 동일하게 설정 후,[0,1) 사이의 값을 가진 momentum coefficient m을 통해 θ_k를 점진적으로 업데이트 합니다. 즉, back-propagation을 통해서는 θq만 업데이트 하는 것이죠. 이를 통해 딕셔너리 안에 있는 key 샘플들은 각자 다른 f_k 에서 발생한 representation이 담기게 되면서, 가장 처음에 enqueue된 샘플들은 제거됨으로써 천천히 일관성을 유지하는 딕셔너리를 구축할 수 있었습니다.

Relations to previous mechanisms

기존 방법들과 비교하여 자신들의 방법론에 대해 간단하게 설명하기도 합니다. (다만 설명이 너무 짧아서 이해하기 위해서는 본 논문을 읽어야,, 사전지식이 있어야,, 이해가 쉬울 듯 합니다..)

먼저 back-propagation에 의한 end-to-end 는 현재의 미니 배치 샘플을 dictionary로 사용하기 때문에 key는 일관되게 인코딩됩니다. 그렇기 때문에 닥셔너리 크기는 mini-batch 사이즈에 연관되어 GPU메모리 크기에 의해 제한되기도 합니다. 게다가 배치가 클 때는 오히려 학습에 어려움이 발생하기도 합니다.

두번째 memory bank는 데이터셋에 있는 모든 샘플의 representation으로 구성되는데요. 각 미니 배치에 대한 딕셔너리는 backpropagation 없이 메모리 뱅크에서 랜덤으로 샘플링되기 때문에 사전의 크기는 크게 가질 수 있습니다. 그러나 representation은 마지막으로 확인되었을 때 업데이트된 것이므로 일관성이 떨어진다는 단점이 있죠.  또한 MoCo가 모든 샘플에 대해 tracking 하지 않는 것에 반해 memory-bank는 메모리 측면상 비효율적이고 데이터가 수억~수십억 개가 넘어갈 경우 구현하기 쉽지 않습니다.

Algorithm: Momentum Contrastive

이제부터는 MoCo 의 알고리즘에 대하여 순차적으로 그림과 함께 설명드리겠습니다.

하나의 이미지에 대하여 서로 다른 Augmentation을 적용하여 Query와 Key 이미지를 생성합니다.

Data Augmentation으로는 랜덤하게 resize된 이미지로부터 224×224 크기로 자른 후, random color jittering, random horizontal flip, 그리고 random grayscale conversion가 적용됩니다.

생성한 쿼리와 키 이미지를 각각 피처를 추출하기 위한 서로 다른 encoder 에 넣어 NxC 크기의 feature를 출력합니다.

앞서 살펴본 바로 Contrastive Learning에는 Positive / Negative Pair가 필요한데요.
MoCo에서는 Dictionary 구조를 사용하여 Self-learning을 진행하기 때문에, 현재 Query 이미지와 가장 유사하게 매칭되는 Dictionary에 저장된 Key 를 찾아야합니다. 매칭되는 Key-Query가 바로 Positive Pair가 되겠죠.

따라서 앞서 같은 이미지로부터 Augmentation된 Query와 Key를 Positive pair로 설정합니다. (그와 반대되는 경우는 negative pair가 되겠죠. ) 그렇기 때문에 앞서 생성한 쿼리 이미지 x_q 와 키 이미지 x_k는 서로 positive pair가 됩니다.

따라서 positive pair 간의 유사도를 계산하기 위해, 앞서 추출한 Feature_Q와 Feature_K에 대하여 matrix multiplication을 진행합니다. 그 결과 Nx1의 벡터가 생성됩니다.

poisitive pair 간 유사도를 계산하였으면 이제 negative pair 간의 유사도를 계산해야 합니다. 앞서 동일한 이미지로부터 augmentation 된 이미지를 제외하고는 모두 negative pair 라고 하였으니, K 크기의 딕셔너리에 저장되어 있는 나머지 Key Feature가 바로 negative pair 가 됩니다.

이제 딕셔너리에 저장되있는 negative pair(KxC) 와 feature_Q(N*C)의 matrix multiplication을 계산하여 NxK의 feature 를 생성합니다. (위 그림에서는 회색 사각형에 해당)

이제 앞서 구한 두 개의 유사도 벡터(positive / negative)를 N x (1+K) 크기의 벡터로 concat 합니다.

N은 배치사이즈인 것을 생각하면 결국 하나의 query에 대해서는 representation (1+K) 크기의 벡터가 생성되게 됩니다.

MoCo에서의 메인 아이디어는 바로 Key와 Query의 representation을 만드는 인코더를 분리하고, 딕셔너리를 구성하는 Key를 생성하는 인코더와 Query를 생성하는 인코더를 서로 다른 속도로 업데이트한다는 것이었습니다.

이를 위해 먼저 바로 전 단계에서 구한 N x (1+K) 벡터를 사용하여 contrastive loss를 계산합니다.

loss 로는 InfoNCE Loss를 사용합니다. InfoNCE Loss는 아래 수식처럼 Positive key와 유사하고 다른 모든 키와는 유사하지 않을 때 값이 작아지도록 설계됩니다. InfoNCE Loss에서 q 는 query를 , k_+ 는 하나의 positive key, k_i는 나머지 K개의 negative key를, \tau{}는 temperature 하이퍼파라미터를 의미합니다. (infoNCE Loss는 직관적으로 query를 positive key로 분류하려는 (K+1) softmax 기반 classification loss 라고 해석할 수도 있다고 하네요)

그리고 구한 Loss에 대하여 query Encoder에 backward 시킵니다.

앞서 query representation을 만드는 인코더와 key representation을 만드는 momentum encoder가 서로 다른 속도로 업데이트 된다고 설명드렸는데요. 논문 제목이 MoCo 인 것도 바로 이번 단계 때문입니다.

그렇다면 왜 이런 수식을 적용하게 되었느냐? 본 논문에서는 queue를 사용하여 dictionary를 크게 확대할 수 있었으나, queue라는 특성으로 인해 backpropagation을 통한 키 인코더 업데이트가 어렵다는 점 때문입니다.

k가 q보다 작은 값으로 업데이트 되도록 Key를 발생시키는 인코더 파라미터에 다음과 같은 Momentum 수식을 적용시킵니다.

m이 바로 0~1 사이의 momentum 계수입니다. momentum은 상대적으로 1에 가까울수록 (0.99~0.999) 더 잘 작동하며, 천천히 업데이트되는 인코더가 큐를 사용하는 핵심임을 확인할 수 있었다고 합니다. 아래 Ablation study를 통해 momentum 계수에 따른 성능 변화를 확인하시기 바랍니다.

이제 마지막으로 구축되어 있는 Dictionary에 계산된 배치 단위의 벡터 (NxC) 를 enqueue 합니다. 그와 동시에 가장 처음에 저장되있던 벡터를 Dequeue 함으로써 Dictionary의 크기는 K로 유지할 수 있게됩니다.

EXPERIMENTS

Dataset MoCo 에서는 1) 1000 개의 클래스에 대해 백만 여개의 이미지를 포함한 ImageNet-1M (IN-1M), 2) 인스타그램의 공개된 1억여개의 이미지로 구성된 Instagram-1B (IG-1B) 를 이용하여 비지도 pre-training을 수행합니다.

Linear Classification Protocol

먼저 feature representation 자체의 성능을 보기 위해, 선형 분류를 진행합니다.

1. 1N-1M에 대하여 unsupervised pretraining 후
2. 인코더를 파라미터를 고정한 뒤
3. supervised linear classifier를 수행합니다. (FC를 추가하여)

먼저 end-to-end/memory bank/MoCo 방식을 비교합니다. 그림을 보면 모든 딕셔너리 크기가 클수록 성능이 좋은 것을 볼 수 있습니다. end-to-end와는 비슷한 성능을 가지는 것을 확인할 수 있는데요, GPU 메모리의 한계로 그래프에는 K즉 배치사이즈가 1024일 때에 대해서만 성능을 보고하였다고 합니다. 즉, 배치 사이즈는 메모리에 제한될 수 밖에 없다는 한계를 가진다고 주장하죠. 그와 더불어 메모리가 충분하더라도 K를 크게할 수록 최적화하는 데에는 어려움이 있고 그에 따라 성능이 보장될지도 미지수라고 합니다. (moco에서의 K는 딕셔너리의 크기입니다)

다른 방법들과의 성능 비교입니다. 가장 기본적인 ResNet-50 (R50)을 이용했을 때 60.6% 로 다른 경쟁 모델에 비해 성능이 더 좋았습니다. 또한 다른 방법들과 마찬가지로 큰 모델일수록 성능이 더 좋습니다.

Transferring features

비지도 학습의 주요 목적은 학습한 feature 를 다른 object detection/segmentation 등의 다른 downstream task에 잘 적용하는 것입니다. 그동안은 ImageNet을 이용한 지도 pre-training 이 메인이었다면, 여기서는 supervised pre-training 과 성능을 비교합니다.

그런데 MoCo와 supervised pre-training과 동일한 설정을 사용하여 실험을 진행하였습니다. 이러한 세팅은 MoCo에 불리해질 수 있지만 MoCo는 경쟁력있는 성능을 보여줍니다.

먼저 PASCAL VOC 데이터에 대한 object detection 성능 결과입니다. Detector는 Faster R-CNN 이며 모든 레이어에 대해 fine-tuning 합니다. R50-dilated-C5 는 ResNet에서 5번째 컨볼루션 스테이지를 2만큼 dilation 시킨 것이며 R50-C4은 backbone이 4번째 컨볼루션 스테이지에서 끝나는 모델입니다. 또한 표3을 통해 다른 contrastive loss 메카니즘과 달리 MoCo의 성능이 더 좋은 것을 확인할 수 있습니다.

---

제안하는 방법은 정말정말 간단한데 성능은 미쳤다(?) 라는 생각을 계속 가지게된 논문입니다.. 역시 저자의 대단함을 또한번 느끼게 되기도 하였네요… 다음에는 더 자세한 설명을 위해 MoCo의 라이벌이라고 할 수 있는 SimCLR에 대해 읽어봐야겠습니다.