Paper :arxiv.org/abs/1911.05722
Momentum Contrast for Unsupervised Visual Representation Learning
We present Momentum Contrast (MoCo) for unsupervised visual representation learning. From a perspective on contrastive learning as dictionary look-up, we build a dynamic dictionary with a queue and a moving-averaged encoder. This enables building a large a
arxiv.org
Code :github.com/facebookresearch/moco
facebookresearch/moco
PyTorch implementation of MoCo: https://arxiv.org/abs/1911.05722 - facebookresearch/moco
github.com
Introduction
MoCo는 Contrasive learning 기반 unsupervised learninig method로서 contrasive learning 적용을 위한 'key' sample을 queue형태로 구성하고, key sample 용도의 encoder를 moving average 형태로 update하는 방식이다.
그래서 아래와 같은 형태를 띈다.
기존의 contrasive learning과 가장 크게 다른 점은 key sample들을 queue형태로 저장한다는 점이다.
그래서 현재의 mini-batch에 있는 sample들은 encoding되어 enqueue되면, oldest는 dequeue된다.
이러한 queue memory를 사용하는 경우, dictionary size를 mini-batch size 이상으로 크게 키울 수 있다는 장점이 있다.
게다가 이전의 key를 저장하고 사용할 수 있기 때문에, key encoder의 변화를 'slowly progressing' 할 수 있다.
이러한 방법을 통해 key encoder가 생성하는 feature의 consistency를 유지할 수 있다.
Contrastive Learning as Dictionary Look-up
Contrasive loss는 다음과 같이 정의된다.
먼저 encoded q와 encoded sample 모음인 dictionary {k0 , k1, k2....} 가 있다고 생각해보자.
이때, dictionary안에 q와 같은 positive sample이 하나 있다고 가정하고 그것을 k+라고 한다.
그리고 나머지는 negative sample이라고 생각하자.
이때 , Loss는 다음 식으로 표현된다.
Momentum Contrast
위의 시각에서 볼때, contrasive learning은 high dimensional space에 discrete dictionary 를 쌓는 과정과 같다.
이때, key가 randomly 하게 sampled 되는 것을 볼때, dictionary 는 'dynamic'하다고 표현할 수 있다.
그리고 key encoder는 training을 거치면서 학습이 되게 된다.
이때, key encoder가 good feature를 학습하기 위해서는 어떻게 해야 될지 고민했고, 다음과 같은 가설을 세우게 된다.
1. good feature는 'large' dictionary로 부터 배울 수 있다. (많은 negative sample이 있으므로, queue를 쓰는 이유)
2. 그러나 encoder를 위한 dictionary key는 가능한 consistent하게 유지될 수 있어야 한다. (encoder가 학습이 진행되더라도)
이러한 가설 아래 저자들은 Momentum Contrast를 제안한다.
Dictionary as a queue.
저자가 제안한 아이디어의 핵심 중 하나는 key 값들을 저장하는 dictionary를 queue로서 구성하는 것이다.
이러한 접근 방식은 이전 mini-batch에서 encoding되었든 key를 다시 사용할 수 있다는 장점이 있다.
즉, dictionary의 size를 minibatch size의 한계에서 벗어나게 하는 장점이 있는 것이다.
또한 이는 hyperparameter로서 조정이 가능하다.
이떄, dictionary에 있는 sample들은 계속해서 update된다.
current minibatch가 dictionary에 enque되면, queue에 있는 oldest mini-batch가 제거된다.
그래서 dictionary는 모든 data의 subset로서 구성된다.
Momentum update
queue를 사용하는 것은 dictionary size를 크게하는 장점이 있지만, back-propagation에 의해서 key encoder를 dictionary 안에 있는 모든 sample에 대해서 계산하여 update하기는 불가능하다는 단점이 있다. (이해 제대로 못했음)
이러한 문제점을 해결하기 위한 간단한 방법은 query encoder를 key encoder로 copy하는 형식으로 update하는 것이다.
그러나 실험 결과가 좋지 않았다.
저자들은 이러한 방식이 key representation의 consistency를 너무 빠르게 변화시키기 때문이라고 가정한다.
그래서 momentum update를 적용한다.
이때, key parameter는 back propagation으로 update되지 않고, query encoder만 update된다.
위 식에 의해서 key encoder는 query encoder에 비해서 더 천천히 update되는 효과를 가져다 준다.
이러한 방법을 통해 queue에 있는 key들은 different encoder에 의해서 encoding이 되지만,
2개의 encoder에 의해서 encoding된 결과의 차이는 줄어들게 된다.
실험결과에 따르면 momentum value가 클수록 더 좋은 결과가 나왔다. (m=0.999 > m=0.9)
Shuffling BN
f_q , f_k 모델 둘다 batch normalization을 사용하는데, 실험결과 BN이 good representation을 학습하는데 오히려 방해가 됨을 확인하였다. BN 과정에서 생기는 batch 내부 안의 정보 공유가 학습을 방해하는 것으로 추측하고 있다.
그리고 저자들은 이 문제를 shuffling BN으로 해결하고 있다.
key encoder f_k 에 대해서 mini-batch의 sample order를 shuffle한다. 그리고 GPU로 전달한다.
그리고 encoding이 끝나면 shuffle 된 것을 다시 되돌린다.
그리고 query encoder f_q의 sample order는 변경하지 않는다.
즉, query와 positive key가 다른 batch상에서 나와야 하며, 같은 batch에 있는 경우 batch normalization 계산과정에서
정보를 공유하는 부분 때문에 더 좋은 representation 학습에 방해가 된다.