[MoCo] Momentum Contrast for Unsupervised Visual Representation Learning (2019) 리뷰

논문 리뷰 동기

2023년 10월 포스코 인공지능연구원에서 컴퓨터 비전 프로젝트(족적 검색 시스템)를 수행하면서 자가 지도 학습(self-supervised learning, 이하 SSL)에 관심을 갖게 되었다. (관련 본인 글: https://seanshkim.tistory.com/114)

MoCo(Momentum Contrast for Unsupervised Visual Representation Learning)는 2019년 Facebook AI Research에서 발표한 논문으로, 컴퓨터 비전 분야의 SSL 방식 모델의 대표 격이라고 할 수 있겠다.

MoCo를 이해한다면 이 모델의 어떤 점을 보완하고 개선하기 위해서 그 뒤로 SimCLR, DINO, MAE 같은 더 발전된 모델이 나왔는지 맥락을 더 쉽게 이해할 것이다.

 

참고로 이 논문의 1저자는 Kaiming He 박사로, ResNet과 He Initialization를 최초로 고안한 분이다.

 

 

 

인트로 Introduction

키워드는 다음과 같다.

• Dictionary, query, key
• Contrastive learning
• Unsupervised Learning

여기서는 unsupervised learning과 self-supervised learning을 같은 뜻으로 혼용한다.

 

사실 자연어 처리 분야에서는 unsupervised learning 패러다임이 아주 성공적이었다.

대표적인 예시로 GPT, BERT 같은 모델은 모두 unsupervised learning 방식으로 학습 자연어 처리 분야에서 성능을 획기적으로 끌어올렸다.

 

그렇지만 computer vision 분야에서는 아직까지도 unsupervised learning보다는 supervised learning 방식이 지배적이다.

왜일까? 그 이유는 “signal(=정보를 말하는 듯) spaces”의 특성 차이 때문이다.

자연어를 다룰 땐 이 신호 공간이 이산적(discrete)이라 할 수 있다. 문장의 단위를 잘개 쪼개다보면 결국 단어와 형태소 등으로 나뉘기 때문이다.

반면에 이미지를 다룰 땐 다루는 정보가 연속적이고 고차원적이고 인간의 소통 방식과는 구조적 차이가 있기 때문에 dictionary를 구축하는 게 힘들다.

그러나 최근 연구에서는 영상 처리 분야에서도 contrastive loss를 이용해서 unsupervised 방식으로 학습을 시도하고 있다.

 

 

 

 

용어 정리: dictionary, key, query

• 여기서 주의!! 용어를 새로 정의하는데, 이런 contrastive loss를 사용하는 방식을 dynamic dictionaries를 짓는 것과 유사하다고 표현한다. 그렇다면 이 논문 저자들이 말하는 dictionary란 무얼 말하는 건지, 왜 dictionary가 dynamic하게 변한다고 표현한 건지 찬찬히 살펴보자.
• dictionary는 원래 key-value 쌍으로 이루어진 자료구조를 말하는데, 이 논문에서 설명하는 key는 이미지나 이미지 조각(patch)에서 추출해서 encoder network에 통과시킨 결과를 말한다.
• Unsupervised learning에서는 query가 key와 유사하도록, key가 아닌 것과는 다르도록(그래서 contrastive-대조되는-loss라는 말을 사용했다) encoder를 훈련시켜야 한다. 이때 encoder란 주어진 이미지를 query로 encoding시키는 encoder를 말한다.
• 그리고 이러한 학습 과정의 기준이 되는 게 바로 contrastive loss이다. 즉 contrastive loss를 줄여나가는 걸 목표로 학습한다.
• 이때 dictionary는 학습하는 과정동안 가능한 1) 크고 2) 일관성이 있어야 바람직하다고 본다. 먼저 dictionary가 커야 하는 이유는 당연하다. 가지고 있는 데이터 샘플이 커야 visual space에 있는 고차원적이며 연속적인 특징을 잘 추출해낼 수 있기 때문이다. 일관성이 있어야 한다는 건 dictionary의 key와 query가 서로 유사한 encoder에 의해 표현되어야 한다는 뜻이다.

 

 

 

Momentum Contrast (MoCo)

1. 데이터 샘플을 queue 형태로 유지한다. 즉, 새로 들어온 데이터(이미지 mini-batch의 encoding된 결과)는 queue에 추가(enqueue)되고 기존에 있었던 오래된 데이터는 queue로부터 삭제(dequeue)되는 방식을 말한다.
2. queue의 이러한 특성 덕분에 dictionary는 mini-batch 크기와 무관하게 큰 규모를 유지할 수 있다.
3. dictionary 안에는 새 key부터 오래된 key가 분포하는데 아무래도 오래된(앞서 만들어진) key가 차지하는 비율이 더 많을 것이다. 그렇기 때문에 key encoder가 너무 급격하게 업데이트되면 안 그래도 큰 규모의 dictionary는 encoder를 제대로 반영하지 못할 수 있다. 그렇기 때문에 encoder는 느리게 변해야 한다(slowly progressing).
4. Pretext task로는 간단한 instance discrimination task가 있다. 즉 똑같은 이미지에 대해 다른 영역을 각각 crop해서 encoding한 결과가 query와 key라면, 서로 유사하기 때문에 매칭을 해야 한다.
5. 이런 pretext task를 활용했더니 ImageNet 데이터셋 기준 linear classification에서 좋은 성능을 보이더라.
6. 결국 unsupervised learning의 목적은 입력 데이터의 representation을 잘 학습하여 downstream task에서도 좋은 성능을 발휘하게 하는 것이다.

 

 

 

 

Method 1: Contrastive Learning as Dictionary Look-up

• 여기서 사용하는 contrastive loss function은 infoNCEloss의 형태로, softmax 함수에다가 log 취하고 - 부호 붙인 꼴과 비슷하다.

 

 

• query q와 key k를 dot product한 결과가 안에 들어가 있다. 람다는 temperature hyperparameter(논문에선 0.07을 사용했다)를 의미한다.
• $k_{+}$라고 밑에 + 기호 붙은 k는 positive key를 의미한다.
• 따라서 query인 q가 key $k_{+}$ 와 유사하고 다른 key k(q 입장에선 이게 negative key이다)와 다르다면, loss function의 값은 낮아야 한다. 그렇기 때문에 이 loss function은 query q를 $k_{+}$ 로 분류(같은 종류로 보겠다는 뜻)하겠다는 (K+1)-way softmax classifier로도 생각할 수 있다.
• 여기서 q는 그냥 샘플 이미지를 말하는 게 아니고 네트워크를 통과해서 encoding된 결과를 말한다. 즉 $q = f_{q}(x^{q})$ 인데, $f_{q}$ 는 encoder network, $x^{q}$ 가 query sample이다.

 

 

 

Method 2: Momentum Contrast

• 아까 dictionary가 dynamic하다고 표현했던 이유는, key들의 모음인 dictionary가 모델 학습 과정 중에도 계속해서 바뀌기 때문이다. 즉 key는 계속 random하게 샘플링하고 key encoder 또한 모델이 훈련을 거치면서 업데이트한다.
• 근데 논문 저자들이 내세운 가설은 이 dictionary가 클수록 좋은 feature를 학습하기 유리할 것이라는 건데, 반면에 key encoder는 가능한 일관성 있게 유지해야 한다는 것이다. 이 조건을 충족하기 위해 momentum contrast라는 개념을 도입했다.

 

Dictionary as a queue

dictionary를 queue의 자료구조로 만든다. queue는 FIFO(First-In-First-Out) 형태의 자료구조다. 현재 막 샘플링한 mini-batch는 dictionary로 들어가고, 가장 오래된 mini-batch부터 queue에서 제거된다. 이렇게 하면 outdated된 key부터 제거할 수 있기 때문에 일관성을 유지하는 데 도움이 된다.

또한 dictionary(queue)의 크기가 mini-batch 크기보다 충분히 크다면 dictionary 내용이 급격하게 변하진 않을 거라 dictionary에 있는 sample은 점진적으로 대체될 수 있다.

 

 

 

 

Momentum Update

애초에 key encoder를 업데이트하려면 back-propagation을 진행해야 하고 그러려면 queue에 있는 모든 sample에 대해 gradient가 역전파되어야 한다.

근데 이건 불가능하니까 처음엔 생각한 건, key encoder gradient 대신 query encoder의 gradient를 그대로 가져와서 업데이트하는 방식이었다. 그러나 이렇게 하면 query encoder의 변화가 key encoder 입장에서 너무나 급격하기 때문에 학습이 제대로 이루어지지 않았다.

momentum m 값을 0.9도 아니고 0.999 정도로 잡으니까 성능이 좋게 나왔다고 한다.

 

 

 

end-to-end vs. memory bank

• end-to-end: 현재 mini-batch가 곧 dictionary인 방식. 그러나 dictionary는 어쨌든 크기가 커야 유리하므로 mini-batch 크기를 늘릴 수밖에 없는데 mini-batch 크기를 무한정 늘릴 순 없다. GPU 처리 용량의 한계가 있기 때문.
• memory-bank: dictionary를 mini-batch 사이즈로 제한하지 말고, 모든 데이터를 dictionary에 다 담아두고 거기서 일부만 샘플링해서 mini-batch로 쓰자. 근데 이렇게 하면 매 스텝마다 전체 dictionary가 아닌, 각 mini-batch의 데이터로만 학습을 하게 된다. 즉 memory bank에 있는 샘플은 이미 예전에 업데이트된 encoder로부터 나온 결과이므로 학습의 일관성이 떨어질 수밖에 없다.

 

 

 

Technical Details & Shuffling BN

• Resnet으로 encoding해서 마지막에 FC Layer 통과한 128-D output vector를 query나 key로 쓴다.
• Batch Normalization이 모델로 하여금 좋은 representation을 배우지 못하도록 막는다. pretext task를 ‘치팅’하는 것처럼 보인다.
• 그래서 multi-GPU 환경을 이용해서, GPU에 mini-batch를 넣기 전 sample 순서를 뒤섞는 과정을 구현한다. 이건 key encoder에만 해당하고 query encoder에 대해서는 진행하지 않는다.