[CVPR 2020] Momentum Contrast for Unsupervised Visual Representation Learning

안녕하세요, 정의철 연구원입니다. 이번에 리뷰할 논문은 2020년 CVPR에 발표된 ‘Momentum Contrast for Unsupervised Visual Representation Learning(MoCo)’입니다. 최근에는 self-supervised learning에 대해 공부하고 있는데, 특히 contrastive learning을 중점적으로 살펴보고 있습니다. 그래서 이번에는 contrastive learning 중에서도 대표적인 논문으로 꼽히는 MoCo를 리뷰하게 되었습니다. 그럼 바로 리뷰 시작하겠습니다.

1. Introduction

자연어 처리에서는 Unsupervised representation learning이 매우 성공적이며, 이를 보여주는 대표적인 예로 GPT와 BERT가 있습니다. 그러나 컴퓨터 비전에서는 supervised pre-training 이 여전히 우세하며, 여기서 unsupervised 방법은 일반적으로 뒤쳐집니다. 이러한 이유는 각각의 signal spaces에 대한 차이에서 비롯될 수 있다고 합니다.

예를 들어 언어 작업은 토큰화된 사전을 구축하기 위한 discrete signal space(단어, 하위 단어 단위 등)을 갖고 있으며, 이러한 특성은 비지도 학습에 유용하게 활용될 수 있습니다. discrete signal space이란 이산적인 신호 공간을 나타냅니다. 이산적인 신호 공간은 개별적인 항목이 분리되어 있으며, 각 항목은 서로 구별됩니다. 예를 들어, 자연어 처리에서 문장을 구성하는 단어들은 이산적인 표현이며 각 단어는 서로 다른 의미를 나타내며 문장 내에서 특정 위치에 배치 됩니다. 반면에 컴퓨터 비전은 연속적인 신호 공간을 사용합니다. 연속적인 신호는 단어와 같이 구조화되어 있지 않기 때문에 처리하기 복잡하고 비지도 학습에 더 큰 어려움을 줄 수 있습니다.

2020년 기준 몇 가지 연구는 unsupervised visual representation learning에 대한 희망적인 결과를 제시했다고 합니다. 이러한 방법들은 대부분 contrastive loss 과 관련된 접근 방식을 사용하여 구축된 dynamic dictionaries이라고 볼 수 있습니다. 사전의 “keys”(토큰)는 데이터(예: 이미지 또는 패치)에서 샘플링되며 인코더 네트워크에 의해 표현됩니다. 비지도 학습은 먼저 인코더를 훈련시켜야합니다. 인코딩된 “query”는 matching 키와 유사하고 다른 키와는 다르도록 만듭니다. 학습은 contrastive loss을 최소화하는 것으로 정의됩니다.

이 관점에서 dictionaries가 가져야할 특성으로 (i) large (ii) consistent 하게 구축하는 것이 바람직하다고 가설을 세웁니다. 직관적으로 큰 dictionary는 continuous, high- dimensional 공간을 더 잘 샘플링할 수 있으며, dictionary의 키는 동일하거나 유사한 인코더에 의해 표현되어야 하므로 쿼리에 대한 비교가 일관되어야 합니다. 저자는 크고 일관된 사전을 구축하기 위한 방법으로 Momentum Contrast(MoCo)를 제안합니다. 저자는 사전을 데이터 샘플의 queue로 유지합니다. 현재 미니배치의 인코딩된 representation이 큐에 인큐되고 가장 오래된 것이 디큐됩니다. 이러한 방식을 통해 더 큰 사전을 가질 수 있도록 합니다. 또한, 키 인코더는 momentum-based moving average 방식을 사용하기에 일관성을 유지할 수 있습니다.

MoCo는 contrastive learning을 위한 dynamic dictionary을 구축하는 메커니즘으로, 다양한 pretext tasks 함께 사용할 수 있습니다. 본 논문에서는 간단한 instance discrimination task으로 사전 학습을 진행합니다. 예를 들어 쿼리와 키가 동일한 이미지에서 인코딩되었다면 일치한다고 판별합니다. 이러한 비지도 학습을 마치면 학습된 피처를 가지고 해결하려는 실제 문제(downstream tasks)에서 fine-tuning하여 비지도 및 지도 표현 학습 사이의 격차를 줄여줍니다.

2. Method

2.1. Contrastive Learning as Dictionary Look-up

먼저 contrastive loss에 대한 설명입니다. contrastive loss는 주어진 쿼리(q)와 dictionary 키들({k0, k1, k2, …}) 간의 유사성을 평가하고 쿼리가 양성 키(k+)에 유사하고 다른 모든 키들과는 다르다는 것을 학습하도록 하는 함수입니다. 수식 (1)은 쿼리(q)를 k+로 분류하려는 K+1개 클래스의 softmax classifier 로그 손실함수로 이해하시면 될 것 같습니다.

결론적으로 위의 contrastive loss는 쿼리와 키를 표현하는 encoder 학습시키는데 사용됩니다.

2.2. Momentum Contrast

Contrastive learning에서는 쿼리와 키 간의 유사성을 비교하여 학습합니다. 따라서 만약 키들이 서로 다른 인코더에 의해 표현된다면, 키 간의 유사성을 올바르게 측정하기 어려울 수 있습니다. 그러나 키들이 동일하거나 유사한 인코더에 의해 표현된다면, 모든 키들은 동일한 표현 공간에 매핑됩니다. 이는 키들 간의 유사성을 일관되게 측정할 수 있게 해줍니다. 이러한 일관성을 유지하기 위해 저자는 Momentum Contrast를 제시합니다.

Dictionary as a queue

접근 방식의 핵심은 큐로 dictionary 유지하는 것입니다. 이를 통해 직전 미니 배치의 인코딩된 키를 재사용할 수 있습니다. dictionary의 샘플은 점진적으로 교체가 되는데 현재 미니 배치가 사전에 인큐되고, 큐에서 가장 오래된 미니 배치가 제거됩니다. 가장 오래된 미니 배치를 제거함으로써 가장 최신 것들과 가장 일관성이 없는 것들은 제거하게 됩니다.

Momentum update

큐를 사용하면 dictionary 크게 될 수 있지만, 키 인코더를 역전파를 통해 업데이트하는데에는 계산적인 어려움이 있습니다. 왜냐하면 기울기는 큐의 모든 샘플로 전파되어야 하는데 큐의 모든 샘플에 대해 그라디언트를 계산하는 것은 실용적이지 않기 때문입니다 . 단순한 해결책은 키 인코더 fk를 쿼리 인코더 fq에서 복사하여 그라디언트를 무시하는 방법이 있지만 이러한 솔루션은 실험에서 좋지 않은 결과를 낳는다고 합니다. 따라서 이 문제를 해결하기 위해 모멘텀 업데이트를 제안합니다.

fk의 매개변수를 θk, fq의 매개변수를 θq로 나타내면, θk는 다음과 같이 업데이트 됩니다.

여기서 m ∈ [0, 1)은 모멘텀 계수이고 매개변수 θq만 역전파로 업데이트됩니다. 식(2)의 모멘텀 업데이트는 θk가 θq보다 더 smoothly하게 업데이트 되도록 만듭니다. 결과적으로, 큐의 키는 서로 다른 인코더(다른 미니 배치에서)에 의해 인코딩되지만, 이러한 인코더 간의 차이를 작게 만들 수 있습니다. 실험에서는 상대적으로 큰 모멘텀(예: m = 0.999, 기본값)이 더 작은 값(예: m = 0.9)보다 훨씬 좋은 결과를 낳으며, 느리게 업데이트 되는 키 인코더가 큐를 사용하는 핵심이라고 저자는 말하고 있습니다.

Relations to previous mechanisms

Fig2에서는 MoCo 포함 3가지 contrastive loss mechanisms을 비교합니다. 세 메커니즘은 키가 유지되는 방식과 키 인코더가 업데이트되는 방식에서 차이가 있습니다. 먼저 (a) end-to-end 메커니즘은 쿼리와 키 표현을 계산하는 인코더는 end-to-end로 업데이트됩니다. (b) memory bank 메커니즘은 키 표현이 메모리 뱅크에서 샘플링됩니다. 메모리 뱅크는 데이터셋의 모든 샘플의 representations으로 구성됩니다. 각 미니 배치의 dictionary는 메모리 뱅크에서 무작위로 샘플링되며 역전파가 없으므로 큰 사이즈 dictionary를 만들 수 있습니다. 그러나 메모리 뱅크의 샘플 representations은 마지막으로 업데이트된 representations을 그대로 사용함으로 representations이 새로운 정보에 적응하는 것을 방해할 수 있습니다.그리고 (c): MoCo는 저자가 제안하는 매커니즘이며 모멘텀으로 업데이트된 인코더에 따라 새로운 키를 동적으로 인코딩하고 키의 큐를 유지합니다

2.3 Pretext Task

Contrastive learning을 위한 Pretext Task에는 다양한 종류가 있습니다. 이 논문에서는 discrimination task를 Pretext Task로 사용합니다.

쿼리와 키가 생성되는 과정은 동일한 이미지에서 쿼리와 키가 생성된 경우를 positive쌍으로 간주하고, 그렇지 않으면 negative 샘플 쌍으로 간주합니다.

이 Pretext Task을 위한 MoCo의 의사 코드는 다음과 같습니다.

1. f_q, f_k 인코더를 초기화 시켜줍니다.
2. 데이터 로터에서 미니배치 사이즈 N만큼 데이터를 가져옵니다.
3. N개의 데이터에 augmentation을 해주어 쿼리와 키를 각각 만들어줍니다.
4. augmentation된 데이터를 각각의 인코더에 넣어줍니다.
5. positive logits과 negative logits을 각각 구해 contrasitve loss를 계산합니다. 이때 negative 샘플은 큐에서 k개 만큼 가져옵니다.
6. 이후 쿼리 인코더와 키 인코더를 업데이트 시켜주는데 이때 키 인코더는 Momentum을 사용하여 업데이트 됩니다.
7. 현재 미니 배치의 데이터는 enqueue, 처음에 있던 미니 배치 데이터는 dequeue하여 dictionary를 업데이트 합니다.

Shuffling BN

저자는 f_q, f_k 인코더에 일반적인 Batch Normalization (BN)을 적용시켰을 때 모델이 좋은 표현을 학습하는 것을 방해한다는 것을 발견했다고 합니다. 이는 모델이 pretext task를 cheating하여 쉽게 낮은 솔루션을 찾기 때문이라고 합니다. 따라서 저자는 이 문제를 해결하기 위해 BN을 섞어 사용합니다. 저자는 여러 GPU로 학습을 수행하고 각 GPU에 대해 샘플에 대해 독립적으로 BN을 수행합니다. 구체적으로 키 인코더 fk의 경우, 현재 미니 배치에서 샘플 순서를 섞은 다음 GPU 간에 분산하기 전에 샘플 순서를 섞습니다. 그리고 쿼리 인코더 fq에 대한 미니 배치의 샘플 순서는 변경되지 않습니다. 이를 통해 쿼리와 키의 positive logits을 계산하는 데 사용되는 배치 통계가 두 가지 다른 하위 집합에서 가져와지도록 보장하여 BN에서의 학습을 가능하게 했다고 합니다.

3. Experiments

다음은 실험 결과입니다. 먼저 데이터 셋으로는 ImageNet-1M (IN-1M), Instagram-1B (IG-1B)을 사용합니다. 인코더에는 ResNet-50모델을 사용합니다.

3.1. Linear Classification Protocol

저자는 먼저 linear classification을 통해 저저가 제안한 방법론을 검증합니다. IN-1M에서 unsupervised pre-training을 수행하고 학습된 features를 가지고 supervised linear classifier 테스크를 수행합니다.

위의 그림은 MoCo 포함 3가지 contrastive loss mechanisms을 비교한 결과입니다.

end-to-end 매커니즘은 K가 작을 때 MoCo와 유사한 성능을 보입니다. 그러나 GPU가 감당할 수 있는 가장 큰 미니 배치는 1024이기 떄문에 K가 높을 때의 비교는 어렵습니다. 또한 K가 더 큰 경우에도 성능이 오를 것이라는 보장은 할 수 없다고 저자는 말하고 있습니다.

memory bank 메커니즘은 더 큰 Dictionary 크기를 지원할 수 있지만 MoCo보다 2.6% 나쁜 성능을 보입니다. 이 결과는 저자가 제시한 가설과 일치하는데 메모리 뱅크는 Memory bank의 sample은 epoch를 돌아서 새롭게 forward가 됐을 때 업데이트된다고 합니다. 이렇게 되면 마찬가지로 memory bank 내부 원소들의 업데이트 주기가 굉장히 길어져서 representation의 일관성을 유지하기가 힘들어집니다. 결과적으로 저자가 제안한 MoCo가 가장 좋은 성능을 보입니다.

momentum에 대한 실험도 진행합니다. 결과적으로 momentum은 0.999일때 가장 좋은 성능을 보이는 것을 확인할 수 있습니다.

위 그림은 기존의 unsupervised learning 모델들과 비교를 한 결과입니다. unsupervised learning 모델들은 모델 크기에 따라 성능 차이가 심할 수 있기 때문에 공정한 비교를 위해 정확도와 모델 파라미터 수를 모두 고려하여 비교를 진행합니다. 또한 저자는 ResNet-50이외에도 채널 수를 늘린 ResNet-50(2×과4×)으로도 성능을 비교하는데 R50w4×로 68.6%의 정확도를 달성하여 가장 좋은 성능을 보여주는 것을 확인할 수 있습니다.

3.2. Transferring Features

비지도 학습의 주요 목표는 중 하나는 전이 가능한 feature를 학습하는 것입니다. 이번 실험에서는 unsupervised 방식으로 사전학습시킨 MoCo 모델과 ImageNet에 대해 supervised 방식으로 사전 학습시킨 모델과의 비교입니다. 비교는 PASCAL VOC , COCO 등과 같은 다양한 작업에서 진행됩니다.

3.2.1 PASCAL VOC Object Detection

먼저 Object Detection의 실험 셋팅에서 detector의 모델로는 Faster R-CNN을 사용하고 backbone모델로는 R50-dilated-C5과 R50-C4를 사용합니다. R50-dilated-C5 는 ResNet에서 5번째 컨볼루션 스테이지에서 2만큼 dilation 시킨 것이고 R50-C4은 backbone이 4번째 컨볼루션 스테이지에서 끝나는 모델입니다.

Table2의 결과표에서 (a)는 Faster R-CNN, R50-dilated-C5의 모델을 사용했을 때 결과, (b)는 Faster R-CNN, R50-C4의 모델을 사용했을 때 결과입니다. 두 결과표를 통해서 알 수 있는 점은 MoCo가 Supervised 방식을 뛰어넘을 수 있다는 점과 detector의 구조에 따라 성능의 차이가 발생할 수 있다는 점입니다. 저자는 pre-training과 detector structures 사이의 관계는 고려해야할 요소 중 하나라고 언급하고 있습니다.

Table3의 결과표에서는 MoCo 포함 3가지 서로 다른 contrastive loss mechanisms을 비교합니다. 보시는 것처럼 MoCo가 가장 좋은 성능을 보임을 확인할 수 있습니다.

Table4에서는 PASCAL VOC에서 fine-tuned된 Object Detection에 대한 기존 pretext task 방법과의 비교 결과입니다. 역시나 기존 pretext task 방법들보다 좋은 성능을 보여주고 있습니다.

4. Conclusion

MoCosms 다양한 컴퓨터 비전 작업과 데이터셋에서 Unsupervised Learning의 가능성을 보여준 논문입니다. 저자는 MoCo의 IN-1M에서 IG-1B로의 개선은 눈에 띄지만 비교적 작으며, 더 큰 규모의 데이터가 완전히 활용되지 않을 수 있다고 말하고 있습니다. 더 큰 규모의 데이터셋을 완전히 활용하기 위해 적절한 pretext tasks의 연구가 필요하다고 저자는 말하고 있는데 비슷한 주제에 대한 최신 논문 및 연구를 찾아보면서 새로운 아이디어와 기술을 정리해 보아야겠네요.

감사합니다.