[PMLR 2020] A Simple Framework for Contrastive Learning of Visual Representations (SimCLR) 핵심 리뷰

1. 들어가며

이번 글에서는 2020년 PMLR에 발표된 A Simple Framework for Contrastive Learning of Visual representations 논문을 리뷰합니다. 이 논문은 SimCLR라고 불리며, 이번 글에서도 SimCLR라고 지칭하겠습니다.

SimCLR는 특히 이미지 분류와 같은 비전 문제에서 뛰어난 성능을 보이며, 라벨이 없는 데이터에서도 유용하게 사용됩니다. 이 글에서는 SimCLR의 기본 아이디어, 그리고 이 방법이 어떻게 다른 자기지도학습 방법들과 차별화되는지에 대해 상세히 알아보겠습니다.

자기지도학습은 라벨이 없는 데이터를 활용하는 방법론으로, SimCLR는 이 중에서도 특히 주목받고 있습니다. 이 글을 통해 SimCLR의 기술적 세부 사항과 그 의미를 이해하고, 실제 파이썬 코드로 어떻게 구현되는지까지 살펴보겠습니다.

2. 제안 방법

SimCLR 이전의 Self Supervised Learning 방법은 대체로 Prediction Task를 학습하는 방식이었습니다. 대표적으로 Context Prediction 방식, Inpainting 방식, Zigsaw Puzzle 방식, Rotation Prediction 방식 등이 있죠. SimCLR는 최초로 Contrastive Learning 방식의 Self Supervised Learning 방법을 제안합니다. Contrastive 는 ‘대조하는’ 이라는 뜻인데요. 이를 통해 이미지를 대조하는 방식으로 학습한다는 것을 유추해볼 수 있습니다. 어떤 방법을 통해, 무엇을 대조하는지 유념하며 SimCLR의 제안 방법을 살펴보겠습니다.

2-1. 핵심 아이디어

먼저 핵심 아이디어를 살펴볼게요. 이미지를 대조하려면 ‘유사’ 한 이미지와 ‘다른’ 이미지가 있어야겠죠? SimCLR에서는 이를 이미지 Augmentation을 통해 만들어 냅니다. ‘유사한’ 이미지와 ‘다른’ 이미지를 모두 Augmentation을 통해 생성해 주는거죠. 그리고 ‘유사한’ 이미지와는 Feature Space에서 가까워지도록, ‘다른’ 이미지는 멀어지도록 학습할 겁니다. 이때 사용해 줄 Augmentation 들은 다음과 같습니다.

SimCLR에서는 Crop, Resize, Flip 그리고 Gaussian Blur 조합의 Augmentation을 사용했습니다.

2-2. Architecture

그래서 전체 구조는 이렇습니다

모델 구조는 Augmentation된 이미지를 압축해주는 Encoder f와 Projection head g로 구성되어 있습니다. 이렇게 나온 Feature z에 대해 ‘유사한’ 이미지면 서로 가까워지도록, ‘다른’ 이미지면 서로 멀어지도록 학습해줄 겁니다.

2-3. 상세 학습 과정

학습 과정을 조금 더 자세히 살펴보겠습니다.

먼저 label 되어 있지 않은 이미지 샘플 N개가 있다고 가정해보겠습니다. 각각의 그림 샘플마다 서로 다른 Augmentation을 2번 수행하여 2개의 Augmented 이미지를 만들어냅니다.

자, 이제 (Augmented) 그림_1 (a) 입장에서 생각해보겠습니다. 이 녀석의 입장에서 보면 (Augmented) 그림_1 (b)는 같은 그림에서 나온 이미지이므로 Sematic 하게 동일한 이미지라고 볼 수 있습니다. 반면 나머지 6개의 이미지는 다른 이미지로부터 파생된 이미지 이므로 Sematic 하게 다른 이미지죠.

즉 위와 같이 구성되게 되는데요. 이때 Positive Pair는 가까워지고, Negetive Pair는 멀어지도록 Loss를 구성해줄 겁니다. 이렇게요.

이때 Sim은 Cosine Similarity를 의미합니다. 아주 간단한 구조죠?

3. 파이썬 구현

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

3-1. Import Module

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

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

3-2. SimCLR Class 정의하기

다음은 SimCLR Class를 정의해줍니다. 사실 SimCLR Class는 매우 단순합니다. 그저 Feature를 추출하는 Encoder와 Projection을 담당하는 Projection Head만이 필요할 뿐이죠.

class SimCLR(nn.Module):
    def __init__(self, base_encoder, projection_dim=128):
        super(SimCLR, self).__init__()
        
        # Base Encoder
        self.base_encoder = base_encoder  # e.g., ResNet
        
        # Projection Head
        self.projection_head = nn.Sequential(
            nn.Linear(self.base_encoder.fc.in_features, self.base_encoder.fc.in_features),
            nn.ReLU(),
            nn.Linear(self.base_encoder.fc.in_features, projection_dim)
        )
        
    def forward(self, x):
        # Feature extraction
        features = self.base_encoder(x)
        
        # Projection
        z = self.projection_head(features)
        
        return F.normalize(z, dim=1, p=2)

3-3. SimCLR 사용하기

이렇게 정의한 SimCLR는 아래와 같이 사용할 수 있습니다. 이 코드에서는 Encoder로 Pretrained ResNet18을 사용하였습니다.

# Example usage
resnet = torch.hub.load('pytorch/vision', 'resnet18', pretrained=False)
simclr_model = SimCLR(base_encoder=resnet)

4. 교훈

다음은 다양한 실험을 통해 SimCLR로부터 얻을 수 있는 교훈에 대해 살펴보겠습니다.

4-1. 첫 번째 교훈

Augmentation 적용 방법에 따른 성능을 보겠습니다.

위 그림은 Augmentation의 구성에 따른 ImageNet Classification 성능을 보여주고 있는데요. x축과 y축이 각각 적용한 Augmentation을 의미합니다. 이 결과로부터 두 가지 교훈을 얻을 수 있습니다.

첫 번째 교훈은 단일 Augmentation 만 적용하면 성능이 매우 안 좋다는 것입니다. 위의 표에서 대각 행렬에 해당하는데요. 단일 Augmentation을 적용하면 대조할 이미지의 개수가 적어지니까 당연한 결과겠죠?

두 번째 교훈은 Crop과 Color Distortion을 같이 구성할 때 성능이 좋다는 것입니다. 이에 대해 저자들은 ‘Crop 만 적용할 경우 Color Distribution이 비슷해서 모델이 이를 가지고 우회 학습이 가능해지는데, 여기에 Color Distortion 까지 추가해주면 이미지의 특징을 학습해야만 하기 때문일 것’ 이라고 추측합니다.

4-2. 두 번째 교훈

Augmentation 정도에 따른 성능을 비교해보겠습니다.

위 표는 Augmentation 정도에 따른 SimCLR 와 Supervised Learning 방식의 성능을 비교한 그림입니다. 위 표를 보면 Color Distortion 이 강해질수록 성능이 좋아지는 SimCLR 에 비해, Supervised 모델은 오히려 성능이 떨어지는 모습을 볼 수 있습니다. 이를 통해 Contrastive Learning 에는 Supervised Learning 에서 보다 더 강한 Data Augmentation 이 필요하다는 교훈을 얻을 수 있습니다.

4-3. 세 번째 교훈

모델의 크기에 따른 성능을 비교해 보겠습니다.

위 그래프는 모델의 크기에 따른 ImageNet 성능을 나타낸 그래프입니다. 빨간색 점은 1000 epoch을 학습한 결과를, 파란색 점은 100 epoch을 학습한 결과를 나타냅니다. 괄호 안의 2x, 4x는 Depth, Width 등의 모델 사이즈를 의미함을 고려하면 모델의 사이즈가 더 클수록 Contrastive Learning으로 학습 했을때의 성능이 더 좋음을 알 수 있습니다.

4-4. 네 번째 교훈

다음은 Projection 방법에 따른 성능을 비교해 보겠습니다.

위 그림은 Projection 방법에 따른 Top1 성능을 나타낸 그래프입니다. Projection Head 를 사용하지 않았을 때보다는 Linear Projection Head 를 사용했을 때가, Linear Projection Head 를 사용했을 때 보다는 Nonlinear Projection Head 를 사용했을 때가 더 성능이 좋은 것을 확인할 수 있습니다.

4-5. 다섯 번째 교훈

다음은 어떠한 Feature를 사용하는것이 더 좋은지 살펴보겠습니다.

위 표는 다양한 Augmentation 에 대해 두 가지 Feature를 사용했을때의 성능을 비교한 표입니다. Representation h는 Encoding 되어 나온 Feature를 의미하고요, g(h)는 이러한 Feature h를 한번 더 Projection 해준 Feature를 의미합니다. 결과를 보면 h를 사용했을 때가 g(h)를 사용했을 때 보다 성능이 높은 모습을 볼 수 있습니다. 저자들은 이에 대해 Transformation 과정에서 정보가 소실되었기 때문일 것으로 추측합니다.

4-6. 여섯 번째 교훈

다음은 다양한 Loss에 대한 성능을 비교해 보겠습니다.

위의 표는 다양한 Loss에 대한 성능을 보여주고 있습니다. 결과를 보면 NT-Xent (Normalized Cross Entropy Loss with Adjustable Temperature)의 성능이 가장 좋은 모습을 볼 수 있습니다.

4-7. 일곱 번째 교훈

다음은 Constrastive Learning 을 사용할때 Similarity 측정 방식에 따른 성능을 비교해 보겠습니다.

위의 표는 L2 Norm과 𝛾에 따른 성능을 보여주고 있습니다. 표를 보시면 L2 Norm을 수행했을 때 (Cosine Simiarity로 측정했을 때) 성능이 더 좋은 모습을 볼 수 있습니다. 또한 𝛾값도 성능에 크게 영향을 주는 모습을 볼 수 있습니다.

4-8. 여덟 번째 교훈

다음은 Batch Size와 학습 시간에 따른 성능을 비교해 보겠습니다.

위의 표는 Batch Size와 Training Epochs에 따른 성능을 나타낸 표입니다. 결과를 보면 Batch Size가 증가할 수록, 오랫동안 학습할 수록 성능이 올라가는 모습을 볼 수 있습니다.

5. 의의

이번 챕터에서는 SimCLR 논문의 의의에 대해 살펴보겠습니다.

첫 번째 의의는 간단하면서도 효과적인 Self Supervised Learning 학습 방법을 제안했다는 점입니다. SimCLR는 자기 지도 학습을 통해 이미지에서 유용한 특성을 추출할 수 있습니다. 이는 레이블이 없는 데이터에서도 효과적으로 학습이 가능하다는 것을 의미합니다. 따라서 레이블링 작업의 비용과 시간을 크게 절약할 수 있습니다.

두 번째 의의는 범용성과 확장성이 좋은 방법을 제안했다는 점입니다. SimCLR는 다양한 아키텍처와 쉽게 통합될 수 있습니다. 기본 Encoder로 어떤 CNN 아키텍처도 사용할 수 있으며, Projection Head도 간단하게 구성됩니다. 이로 인해 SimCLR는 다양한 문제와 데이터셋에 적용할 수 있어 범용성과 확장성이 뛰어납니다.

이 두 가지 의의는 SimCLR가 단순히 이미지 특성을 추출하는 것을 넘어, 레이블이 없는 데이터에서도 높은 성능을 달성할 수 있는 강력한 프레임워크라는 것을 보여줍니다. 이는 자기 지도 학습의 가능성을 크게 확장하는 중요한 발전입니다.

6. 마치며

이번 글에서는 자기지도학습의 한 형태인 SimCLR에 대해 깊게 다루었습니다. SimCLR의 핵심 아이디어와 그 특징을 이해하는 데 중점을 뒀으며, 실제 파이썬 코드를 통해 구현 방법도 살펴보았습니다.

SimCLR의 의의에 대해서도 논의했습니다. 이 방법은 라벨이 없는 데이터에서도 효과적으로 특징을 추출할 수 있어, 다양한 비전 문제에서 뛰어난 성능을 보이고 있습니다. 이는 자기지도학습의 가능성을 크게 확장하는 중요한 점입니다.

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

Ffightingseok