[논문 리뷰] MoCo: Momentum Contrast for Unsupervised Visual Representation Learning(2020)

이번에 리뷰할 논문은 Momentum Contrast for Unsupervised Visual Representation Learning 입니다.

 

https://paperswithcode.com/paper/momentum-contrast-for-unsupervised-visual

Papers with Code - Momentum Contrast for Unsupervised Visual Representation Learning

 

---

본 포스팅을 이해하기 위해서는 "contrast learning" 의 기본적인 이해가 필요합니다. 
혹여나 처음 보는 개념이라면 contrast learning의 구조와 진행 방식 정도를 간단히 이해하고 보시는 것을 추천드립니다.

---

 

 

 

 

 

 

이번 논문에서 제안하는 Momentum Contrast (MoCo) 는 self-supervised learning 중 contrast learning 의 한 종류입니다. 해당 기법은 queue를 사용하여 dynamic dictionary를 구축한다는 특징이 있죠. 

 

NLP task의 경우 words, sub-word units 와 같이 discrete signal space를 가지기 때문에  unsupervised learning을 사용하기 위한 토큰화된 dictionary를 구축하기 수월합니다. 하지만 Computer Vision의 경우 raw signal이 연속적이고, 고차원 공간에 존재하기 때문에 구조화되어있지 않아 dictionary를 구축하는데 어려움이 있습니다.

 

이런 어려움을 해결하기 위해 나온 기법들 중 하나가 바로 contrast learning 입니다. 이는 일반적으로 image data(patch) 에서 sampling 한 후 encoder network를 통해 추출된 representation을 key로 정의한 후 query와 유사도를 비교하는 식으로 진행됩니다. 

 

 

 

 

본 논문은 한 가지 가설을 세웠습니다.

 

 

dictionary를 large 하고, consistent 하게 구축하는 것이 학습을 하기에 바람직하다.

 

 

결국 dictionary를 구축하는데 있어 고려해야할 점은 "large" 와 "consistent" 인데 일반적으로 큰 dictionary가 연속적이고, 고차원의 visual space를 더 잘 sampling 할 수 있겠죠. 또한 positive pair와 negative pair의 representation을 비교하는 방법인만큼 dictionary의 key들을 같거나 비슷한 encoder로 represent 하도록 하여 query의 비교가 consistent 하게 유지해야합니다. 

 

하지만 뒤에서 소개해드릴 contrast loss를 사용하는 기존 기법들은 두 가지 조건 중 하나에서 제한적이었다는 한계가 존재했습니다. 본 논문은 이 large하고 consistent 한 dictionary를 구축하기 위해서 Momentum Contrast (MoCo) 를 제안하게 됩니다.

 

 

MoCo 는 각 mini-batch의 encoding된 image들이 queue에 추가되며, 일정 크기를 설정하고 오래된 이미지 순으로 dequeue하는 방식으로 queue를 유지합니다. 이를 통해 mini-batch의 크기와 상관없이 독립적으로 queue의 크기를 크게 유지할 수 있으며, 더 다양한 데이터를 유지하고 학습에 활용하는 것이 가능해집니다. 

 

또한 key encoder를 query encoder의 momentum-based moving average를 사용하여 천천히 update를 진행합니다. 즉, 사전의 key들을 천천히 update하게 되는데 이를 통해 각 mini-batch 사이에 dictionary의 consistent를 유지할 수 있게 됩니다.

 

 

 

 

 

 

 

 

 

Contrastive Loss - InfoNCE

contrative loss를 구하기 위해서는 우선 query와 key를 sampling 해야합니다. 그 후에는 각 sample들을 encoder network를 통해 representation을 생성합니다. sampling한 query와 key를 $ x^{q} $ 와 $ x^{k} $ 라고 했을 때 encoding 과정을 식으로 표현하면 다음과 같습니다.

 

• $ q = f_{q}(x^{q}) $     
• $ p = f_{p}(x^{p}) $

 

각 key, query의 representation을 생성했으니 이제 contrastive loss 를 정의해야합니다. 본 논문에서는 contrastive loss를 dot product로 유사도를 측정하는 InfoNCE 함수를 사용합니다.

 

• $ q $ : encoded query
• $ \left\{ k_{0}, k_{1}, k_{2}, ... \right\}  $ : keys of dictionary
• $ k_{+} $ : positive key

 

contrastive loss는 $ q $ 가 positive key인 $ k_{+} $ 와는 유사하고, negative key인 나머지 key들에 대해서는 불일치할 때 가장 값이 minimize 됩니다. query와 key의 representation 유사도는 dot product를 수행하여 진행되죠.

 

위 함수의 분자는 query와 positive pair인 key 의 representation 유사도를 측정하며, 분모의 경우 1개의 positive key와 K개의 negative key의 유사도 합으로 구성됩니다. 

 

 

 

 

 

 

 

 

 

Momentum Contrast

이미지의 경우 데이터가 연속적이며 고차원이기 때문에 이를 고려한 discrete dictionary를 구축해야합니다. 위에서 언급한대로 본 논문은 좋은 feature를 추출하기 위해 dictionary를 large하고 consistent 해야한다고 했었죠. 이를 위해 Momentum Contrast 방법론을 제안합니다. 

 

 

 

Dictionary as a queue

queue를 사용하면 dictionary의 크기를 임의로 조절할 수 있기 때문에 크게 유지하는 것이 가능해집니다. 또한 오래된 mini-batch의 encoding된 key들은 최신의 key들과 consistent가 떨어지기 때문에 오히려 제거하는 것이 바람직하다고 합니다. 

 

하지만 위와 같이 queue를 사용하여 dictionary를 구축하는 것은 key encoder를 back propagation으로 update하는 것을 비효율적으로 만듭니다. gradient가 queue에 있는 모든 sample들로 전달되어야 하기 때문이죠. 즉, key의 encoder를 update하는 과정에서 문제가 생기기 때문에 이를 momentum 을 사용하여 update를 시도합니다.

 

• $ \theta_{q} $ : parameters of Query
• $ \theta_{k} $ : parameters of Key
• $ m \in \left [ 0,1 \right ) $ : momentum coefficient

 

$ \theta_{q} $ 만 back propagation 으로 update를 수행하며, $ \theta_{k} $ 는 momentum을 주어 $ \theta_{q} $ 에 비해 훨씬 더 smooth 하게 update 됩니다. 결국 dictionary가 consistent 하지 못하게 되는 이유는 key 에 적용되는 encoder가 단계적으로 update되기 때문에 서로 다른 representation을 생성하기 때문입니다. 물론 momentum contrast를 사용하더라도 mini-batch마다 사용되는 encoder는 달라지지만 encoder 사이의 차이를 작게 유지할 수 있기 때문에 consistent를 유지할 수 있게 됩니다. 추가로 mini-batch 마다 새로 update를 진행하기 때문에 효율적으로 학습하는 것이 가능해집니다.

 

본 논문에서는 $ m $ 이 클수록 잘 작동한다고 합니다. 즉, query 의 영향을 조금씩 받으면서 천천히 update하도록 유도하는 것이죠. 이는 천천히 진화되는 key encoder가 queue를 활용하는 핵심이라 볼 수 있습니다. 저자는 $ m $ 을 0.999 로 설정하여 실험을 진행했다고 합니다.

 

 

 

Relation to previous mechanisms

자 그러면 MoCo mechanism 이 이전의 일반적으로 사용되었던 mechanisms 들과 비교해보도록 합시다.

 

 

 

 

< end-to-end >

해당 기법은 현재 mini-batch의 sample을 dictionary로 사용합니다. query와 key가 일관되게 sampling되기 때문에 key들은 consistent하게 encoding 됩니다. 

 

하지만 dictionary의 크기는 오로지 GPU memory 크기에 제한되는 mini-batch 의 크기로 정해지기 때문에 large dictionary를 구성하는데에는 제한적입니다. 또한 batch size를 너무 크게 가져가면 학습에 문제가 발생할 수 도 있죠.

 

• Large : X
• Consistent : O

 

 

< memory bank >

memory bank는 모든 dataset의 모든 sample들의 representation 으로 구성됩니다. 각 mini-batch의 dictionary는 무작위로 sampling 되기 때문에 각 mini-batch에 대해 key encoder에는 back propagation는 수행되지 않습니다. 이 방법을 사용하게 되면 모든 key들의 representation을 dictionary로 구성할 수 있기 때문에 large dictionary를 구성할 수 있습니다.

 

하지만 각 sample의 embedding이 서로 다른 시점에서 갱신될 수 있기 때문에 consistent함이 부족해지는 결과를 초래합니다.

 

• Large : O
• Consistent : X

 

 

< MoCo >

MoCo는 momentum update를 사용하기 때문에 모든 sample들을 추적하지 않습니다. 이로인해 효율적으로 학습하는 것이 가능하며 encoder 간의 차이를 작게 유지하기 때문에 consistent하게 유지하는 것이 가능해집니다.

 

또한 queue를 사용하여 dictionary를 구축하기 때문에 large dictionary를 구성하는 것이 가능해집니다.

 

• Large : O
• Consistent : O

 

 

위 3가지 기법들을 비교하여 실험한 결과 다음과 같은 결과를 확인할 수 있습니다.

 

 

 

 

 

 

 

 

 

 

Technical details

본 논문에서는 encoder로 ResNet을 사용합니다. ResNet의 convolution 연산과 global average pooling을 적용한 후 fully-connected layer를 통해 생성되는 output의 차원을 128-D 로 고정시킵니다. 

 

이 output vector에 대해 L2-norm으로 normalize를 진행한 후 이를 query 혹은 key의 representation으로 정의합니다. 

 

다음으로 각각 data augmentation을 거치게 되는데 기본적으로 224x224 로 crop을 진행한 후 augmentation을 적용합니다. 적용한 각 기법들은 다음과 같습니다.

 

• randomly resized image
• random color jittering
• random horizontal flip
• random grayscale conversion

 

 

 

 

 

 

 

 

 

 

Pseudocode of MoCo

다음은 전체적인 MoCo 학습 과정을 pseucode로 구현한 것을 확인할 수 있습니다.