[CVPR 2020] Momentum Contrast for Unsupervised Visual Representation Learning (MoCo) 핵심 리뷰

1. 들어가며

이번 글에서는 2020년 CVPR에 발표된 Momentum Contrast for Unsupervised Visual Representation Learning 논문을 리뷰합니다. 이 논문은 MoCo라고 불리며, 이번 글에서도 MoCo라고 지칭하겠습니다.

MoCo는 특히 이미지 분류, 객체 탐지 등 다양한 비전 문제에서 뛰어난 성능을 보이고 있습니다. 이 글에서는 MoCo의 핵심 원리와 이 방법이 어떻게 다른 자기지도학습 방법들과 차별화되는지에 대해 깊게 다루겠습니다.

자기지도학습은 라벨이 없는 데이터를 효율적으로 활용하는 방법론입니다. MoCo는 이 중에서도 ‘모멘텀’을 이용해 효과적인 특징 추출을 가능하게 합니다. 이 글을 통해 MoCo의 기술적 세부 사항과 그 의미를 이해하고, 실제 파이썬 코드로 어떻게 구현되는지까지 살펴보겠습니다.

2. 기존 방법의 문제점

MoCo 전에도 다양한 Self Supervised Learning 방식이 있었는데요. 이중에서도 Contrastive Learning을 사용하는 방법들과 비교해보면 좋습니다.

2-1. End to End 방식

먼저 End to End 방식이 있습니다. 대표적으로 SimCLR를 들 수 있는데요. Self Supervised Learning 하면 빠질 수 없는 방법이죠.

하지만 End to End 방식에는 치명적인 단점이 있습니다. 바로 Batch Size에 크게 종속적이라는 점입니다. 무슨 말인지 조금 더 생각해 볼게요. Contrastive Learning 은 Positive Pair 는 당기고 Negative Pair 는 밀어내도록 학습하는 방식이잖아요. 이때 아주 많은 Negative Sample이 필요합니다. 왜냐하면 Negative Sample들은 밀어내는 방향을 정해주는 역할을 하는데, 고차원상에서 일관된 방향을 알려주려면 아주 많은 샘플이 필요하기 때문이죠. 따라서 매우 큰 Batch Size 를 요구하는데요. GPU 메모리의 한계가 있으니 무한정 Batch Size 를 늘려줄 수는 없다는 한계가 있습니다.

2-2. Memory Bank

그래서 이러한 한계를 극복하기 위해 제안된 방법이 Memory Bank 방식입니다.

위의 논리를 따라간다면 다음에 나올 아이디어는 당연히 아주 많은 샘플을 저장해놓고 비교하면서 학습하자는 것이겠죠? 이러한 Memory Bank 방식은 End to End 방식의 한계를 극복할 수 있었습니다. 하지만 Memory Bank의 많은 샘플들 중에서 Sampling을 통해 학습을 할 텐데요. 이렇게 Sampling을 할 때마다 밀어내야 하는 방향이 바뀐다는 단점이 존재했습니다. 즉 학습이 Consistent 하지 못하다는 것이죠.

3. 제안 방법

그래서 MoCo는 어떻게 하자는 걸까요?

3-1. 핵심 아이디어

위 흐름 그대로 MoCo의 핵심 아이디어를 살펴보겠습니다.

Memory Bank의 한계는 학습이 일관되지 않다는 거였죠? 그럼 당연히 MoCo 에서는 일관된 학습을 할 수 있는 방법을 제안하겠죠? 그 방법이 바로 Momemtum을 이용한 학습입니다. MoCo의 이름에 Mo는 Momentum을, Co는 Contrastive를 의미하죠. 위 그림에서처럼 Encoder를 바로바로 업데이트하는게 아닌, Momentum을 사용해서 천천히 업데이트 해주는 겁니다. 조금더 자세히 살펴보죠.

3-2. Architecture

MoCo의 전체 Architecture는 이렇습니다.

SimCLR와 비교하면서 생각해보면 쉬운데요. 가장 큰 차이는 Encoder가 2개가 생겼다는 것이죠. SimCLR에서는 한개의 Encoder만 사용했는데요. MoCo에서는 2개의 Encoder를 사용합니다. 각각 Q Encoder, K Encoder 라고 부를게요. Q Encoder는 기존과 동일하게 Contrastive Loss로 바로바로 업데이트 해줍니다. 그럼 Feature 분포가 급격하게 변하니까 학습이 Consistent하지 못하겠죠? 따라서 K Encoder는 바로바로 업데이트해 주지 않습니다. 대신 Q Encoder를 천천히 따라가는거죠.

3-3. Momentum Update

이 부분을 Momentum을 사용한 Update 방법이라고 부르는데요. 말은 거창하지만 수식으로 보면 간단합니다.

K Encoder의 가중치 𝜃_k는 Q Encoder 가중치 𝜃_q와 𝜃_k의 가중 평균으로 업데이트되고 있죠. 이때 m은 0.999 정도의 값을 사용합니다. 이 말은 𝜃_q의 값은 0.001 만큼만 반영하겠다는 의미죠. 즉 아주 아주 천천히 업데이트하겠다는 의미입니다. 이 방법을 통해서 Consistent한 학습을 할 수 있었던 것이죠.

3-4. Memory Bank

Memory Bank 부분도 짚고 넘어가지 않을 수 없는데요. SimCLR와 달리 MoCo에서는 아주 많은 샘플을 사용하기 위해 Memory Bank를 구성했죠. 이때 Memory Bank에 샘플이 들어가고 나오는 방식은 Queue 형태를 사용합니다. 먼저 들어간 샘플이 먼저 나오는 방식이죠. 그래서 Memory Bank 안에는 들어온 순서대로 샘플이 차곡차곡 쌓여 학습에 사용되는 방식입니다.

3-5. Loss

Loss는 기존 방법과 동일하게 infoNCE Loss를 사용했습니다.

Positive Sample과 가까워지도록, Negative Sample과 멀어지도록 Q Encoder를 업데이트하게 됩니다.

4. 파이썬 구현

이번 챕터에서는 파이썬을 사용하여 MOCO를 직접 구현해봅니다. 이 과정을 통해 위에서 살펴본 MOCO의 내용을 정확하게 이해해봅니다.

4-1. Import Module

먼저 필요한 Module들을 Import 합니다.

import torch
import torch.nn as nn
import torch.nn.functional as F

4-2. MoCo Class 정의하기

이어서 MoCo Class를 정의해줍니다. MoCo의 핵심은 Momentum Update 방법인데요. __init__ 함수에서 Momentum Update 방법을 설정해놓고, forward 함수에서 Momentum Update를 구현한 모습을 볼 수 있습니다.

class MoCo(nn.Module):
    def __init__(self, base_encoder, dim=128, K=4096, m=0.99):
        super(MoCo, self).__init__()
        
        # Base Encoder
        self.encoder_q = nn.Sequential(
            base_encoder,
            nn.Linear(base_encoder.fc.in_features, dim)
        )
        self.encoder_k = nn.Sequential(
            base_encoder,
            nn.Linear(base_encoder.fc.in_features, dim)
        )
        
        # Momentum update
        for param_q, param_k in zip(self.encoder_q.parameters(), self.encoder_k.parameters()):
            param_k.data.copy_(param_q.data)
            param_k.requires_grad = False
        
        # Queue
        self.register_buffer("queue", torch.randn(dim, K))
        self.queue = nn.functional.normalize(self.queue, dim=0)
        
        self.K = K
        self.m = m

    def forward(self, x_q, x_k):
        q = self.encoder_q(x_q)
        k = self.encoder_k(x_k)
        
        # Momentum update for key encoder
        with torch.no_grad():
            for param_q, param_k in zip(self.encoder_q.parameters(), self.encoder_k.parameters()):
                param_k.data = param_k.data * self.m + param_q.data * (1. - self.m)
        
        return q, k

4-3. MoCo 사용하기

이렇게 정의한 MoCo Class는 이렇게 사용할 수 있습니다.

resnet = torch.hub.load('pytorch/vision', 'resnet18', pretrained=False)
moco_model = MoCo(base_encoder=resnet)

5. 실험 결과

다음은 다양한 실험 결과를 통해 MoCo의 성능을 확인해 보겠습니다.

5-1. Linear Classification

먼저 Linear Classification 성능입니다. 이 결과를 통해 MoCo가 얼마나 유용한 Feature를 학습했는지 알 수 있죠.

위 표는 기존 방법들과 비교한 Linear Classification 성능입니다. 기존 방법들중 MoCo가 가장 좋은 성능을 보이는 모습을 볼 수 있습니다.

5-2. Transferring Features

다음은 다양한 Downstream Task에 대한 Transfer Learning 성능을 확인해 보겠습니다.

먼저 Object Detection 성능입니다.

CoCo Detection과 Instance Segmentation 성능입니다.

기타 다양한 Downstream Task 성능입니다.

모두 기존 방법들보다 성능이 좋음을 알 수 있습니다.

6. 의의

다음은 MoCo 논문의 의의를 정리해보겠습니다.

첫 번째 의의는 Contrastive Learning 방법의 효율성을 향상시켰다는 점입니다. MOCO는 모멘텀 업데이트와 큐(queue)를 사용하여 대조적 학습을 효율적으로 수행합니다. 이를 통해 레이블이 없는 데이터에서도 높은 성능의 특성을 추출할 수 있습니다. 모멘텀 업데이트는 키 인코더의 파라미터를 부드럽게 업데이트하므로, 학습이 더 안정적이고 효율적입니다.

두 번째 의의는 레이블이 없는 데이터를 효과적으로 활용할 수 있는 방법을 제안했다는 점입니다. MOCO는 자기 지도 학습을 통해 레이블이 없는 데이터에서도 유용한 특성을 학습할 수 있습니다. 이는 레이블링에 드는 비용과 시간을 절약하면서도, 높은 성능의 모델을 구축할 수 있음을 의미합니다.

이 두 가지 의의는 MOCO가 단순히 레이블 없는 데이터에서 특성을 추출하는 것을 넘어, 대조적 학습의 효율성과 유용성을 크게 향상시키는 중요한 기술적 발전이라고 할 수 있습니다. 이로 인해 MOCO는 자기 지도 학습 분야에서 매우 중요한 위치를 차지하고 있습니다.

7. 마치며

이번 글에서는 자기지도학습의 중요한 방법 중 하나인 MOCO에 대해 깊게 살펴보았습니다. MOCO의 핵심 원리와 그 특징, 그리고 실제 파이썬 코드를 통한 구현 방법에 대해 알아보았습니다.

MOCO의 의의에 대해서도 논의했습니다. 이 방법은 ‘모멘텀’을 이용하여 라벨이 없는 데이터에서도 효과적으로 특징을 추출할 수 있습니다. 이로 인해 MOCO는 다양한 비전 문제에서 뛰어난 성능을 보이고 있으며, 자기지도학습의 가능성을 더욱 확장하고 있습니다.

마지막으로, MOCO와 자기지도학습은 딥러닝의 미래에 큰 영향을 미칠 것으로 보입니다. 이 글을 통해 MOCO의 중요성과 그 기술적 세부 사항에 대한 깊은 이해를 얻으셨기를 바랍니다. 다음 글에서는 또 다른 흥미로운 딥러닝 아키텍처와 기술에 대해 알아보겠습니다.

Ffightingseok