[NeurIPS 2020] Bootstrap Your Own Latent A New Approach to Self-Supervised Learning

안녕하세요 정의철 연구원입니다. 제가 이번에 소개할 논문은 ‘Bootstrap Your Own Latent A New Approach to Self-Supervised Learning’이고 흔히 BYOL이라고 불리는 논문입니다. 저는 최근 Self-supervised learning 분야에 대해 공부를 하고 있었습니다. 그리고 지속적으로 최신 논문 위주로 follow up 하고 있었습니다. 하지만 Introduction이나 related work에 항상 언급되는 BYOL을 보면서 그냥 넘어가기에는 이후에 나온 논문들의 이해도가 떨어질 것 같아 이렇게 리뷰하게 되었습니다. 그럼 바로 리뷰 시작하겠습니다.

먼저 BYOL에 대해 알아보기 전에 Self-supervised learning은 어떤 task인지 간단하게 알아보겠습니다.

Self-supervised learning

먼저 배경 Self-Supervised Learning이 나오게된 배경은 일반적으로 Supervised Learning 모델을 훈련시키는게 유리하지만 수많은 데이터에 label을 다는 것에는 한계가 있습니다. 따라서 이에 대한 방안으로 unlabelled dataset으로부터 좋은 representation을 얻고자하는 학습방식을 고안하게 되었고 이를 Unsupervised Learning이라 부릅니다. 그리고 Unsupervised에서 unlabelled dataset으로부터 좋은 representation을 뽑는 task를 self로 정하고 모델을 학습하는 것을 사전학습이라고 부릅니다. 그리고 이렇게 학습된 representation을 이용하여 다른 task에 적용시키는 것을 Self-Supervised Learning이라고 부릅니다. 여기서 언급한 사전학습에는 크게 두 가지가 있습니다. 그 종류에는 pretext task , contrastive learning이 있습니다.

pretext task

pretext task에는 다양한 방법들이 있지만 그 중에서도 rotation task를 예시로들어 설명드리겠습니다.

위 그림에서 첫 번째 Stage가 사전학습을 진행하는 부분입니다. 먼저 이미지 X에대해 0,90,180,270도로 회전을 시켜 라벨을 붙여줍니다. 그리고 이 데이터 셋으로 모델 f를 학습시키게 됩니다. 이후 두 번째 Stage에서는 이렇게 학습된 f를 freeze시키고 뒤에 선형 레이어만 붙여서 학습시키는 모습입니다.

contrastive learning

한편 Contrastive learning은 별도의 pretext task가 주어지지 않고 대신에 encoder에서 나온 representation을 사용하여 사전학습을 진행하게 됩니다. 이를 위해서는Anchor, Positive pair , Negative pair가 필요합니다.

위 그림은 Contrastive learning의 대표 모델인 SimCLR의 구조도 입니다. 먼저 원본이미지 (Anchor)에서 augmentation시켜 이미지 두 개를 생성하는데 이 1쌍을 Positive pair라고 부릅니다. 그리고 다른 Anchor에서 augmentation시킨 이미지들과는 Negative pair 관계를 구성하게 됩니다. 이렇게 데이터를 만든 뒤에는 Encoder에 넣어 Representation 벡터를 뽑고 이를 Projection head에 넣어 차원이 간소화된 벡터 Z를 뽑게 됩니다. 이후 앞에서 구성한 Positive pair , Negative pair 관계에 따라 Positive관계의 벡터는 embedding space에서 거리가 가깝게 Negative관계는 거리가 멀게 하도록하는 Contrastive Loss를 계산하여 학습이 진행됩니다.

Contrastive learning의 단점

하지만 이러한 Contrastive learning에도 단점이 존재하는데 다음과 같습니다.

1. Negative pair를 제공하기 위한 전략필요

• Contrastive learning에서 충분한 성능을 달성하기 위해서는 많은 양의 Negative pair와의 비교가 필요합니다. 그리고 이러한 Negative pair를 구성하기 위한 전략이 필요합니다. 예를들어 SimCLR는 배치 사이즈에 크게 영향을 받기 때문에 Large Batch를 사용해야만 성능이 나오고 이를 개선하기 위해 MoCo라는 방법론은 Memory Queue를 도입하여 별도의 공간에 Negative representation을 저장해줍니다.

2. Choice of Image Augmentation

• Contrastive learning의 두번째 단점으로는 Augmentation 조합에 따라 모델의 성능이 크게 좌우된다는 것입니다.

위의 그림은 SimCLR 논문에 나온 Augmentation 조합에 따른 성능표 입니다. 보시면 Color distortion + Crop 조합에서는 가장 좋은 성능을 달성했지만(55.8) 다른 조합에서는 성능이 많이 떨어지는 것을 확인할 수 있습니다.

3. Requires many negative samples

• 세 번째 단점으로는 이전에 설명드렸듯이 Contrastive learning으로 충분한 성능을 달성하기 위해 수 많은 negative samples이 필요하다는 것입니다.

Motivation BYOL

지금까지 전반적인 Self-supervised Learning에 대해 설명드렸고 그 중에서도 Contrastive learning을 중점적으로 설명드렸습니다. 그리고 Contrastive learning에는 Negative pair를 셋팅하고 활용하는데 약간의 어려움이 있었기에 본 논문에서는 Negative pair를 사용하지않고 Positive pairs만 사용하여 기존 Contrastive Learning 방식들보다 우수한 성능을 달성합니다.

위의 그림은 BYOL에 영감을 주었던 간단한 실험입니다. 이 실험은 positive pair만을 사용하여 좋은 레프센테이션을 학습할 수 있는지에대한 실험입니다. 먼저 랜덤 초기화된 모델 A의 파라미터를 고정하고 그뒤에 선형레이어를 하나 쌓아서 모델을 학습시키고 성능을 측정하게 됩니다. 당연하게도 성능은 좋지 않고 Top1 ACC는 1.4퍼센트를 찍게 됩니다. 다음 실험으로는 위에 있는 모델 구조와 모두 동일한 모델 B를 만들어 줍니다. 하지만 이때 모델 B 뒤의 선형레이어는 라벨 Y로 학습하는게 아니라 이전에 모델 A가 예측했던 Y’ 값으로 학습을 진행하게 됩니다. 이를 데이터셋으로 하면 모델 B도 성능이 않좋게 나올 것 같지만 성능을 확인해보면 Top1 ACC는 18.8퍼센트를 찍게 됩니다. 여전히 top1 어큐러씨가 낮기는 하지만 그래도 A모델 보다는 레프레센테이션을 잘 학습했다고 볼 수 있고 이 결과가 BYOL을 만드는 계기가 되었다고 합니다.

Method

BYOL의 아키텍처는 다음과 같습니다. 먼저 2개의 네트워크로 구성이 되어있고 윗 부분이 온라인 네트워크에 해당하고 아랫부분이 타겟 네트워크에 해당합니다. 이미지가 주어졌을 때 모델이 작동하는 방식을 살펴보면 먼저 하나의 이미지에대해 각기 다른 augmentaion를 하고 이 뷰들을 각각의 인코더에 입력으로 넣어 representation y를 뽑습니다. 또 이 값을 projector에 넣어 projection된 벡터 z를 뽑게 됩니다. 근데 이때 BYOL의 학습은 온라인에서 나온 벡터z를 predictor에 넣어 타겟 벡터 z’을 예측하도록 학습이 진행됩니다. 그리고 또 하나 차이점은 이 부분에서 예측을 하고 나온 로스값으로 온라인 네트워크의 경우 파라미터를 업데이트 하지만 타겟 네트워크의 경우 그레디언트를 업데이트 하지 않고 온라인 네트워크의 파라미터 값을 가중 평균한 값으로 업데이트를 하게 됩니다. 이렇게 타겟 네트워크가 파라미터를 업데이트 하는 방식을 Exponential Moving Average라고 합니다.

위의 식이 타겟 네트워크의 파라미터를 업데이트하는 exponential moving average 공식입니다. 이때 τ의 값은 0.996에서 시작해 cosine annealing을 적용시켜 iteration이 진행될수록 0.996부터 시작해서 점점 1에 가까운 값으로 키워주게됩니다.

또한 loss function으로는 mean squared error function을 사용하였고, online network와 target network의 projection output들에 l2 normalization 시킨 값들 간의 mean squred error를 구하여 loss를 계산하게 됩니다.

마지막으로 학습을 대칭적으로하기 위해 augmentaion 한 데이터를 교환해서 한 번더 loss 값을 구한 뒤 이 두 loss 값을 합쳐서 최종 로스값을 만들어 내게 됩니다.

Experimental evaluation

다음은 실험 부분입니다.

우선 Table 1에 대한 실험은 linear evaluation에 대한 실험으로 ImageNet unlabeled 데이터셋으로 Self Supervised Learning 방식으로 Pre-training을 시킨 뒤, feature extractor(encoder)를 freeze 시키고, linear classifier를 학습시키는 linear evaluation하는 실험입니다. 실험 결과를 확인해보면 SimCLR, MoCov2와 같은 Contrastive Learning 방법론들보다 더 좋은 성능을 달성하였습니다. 그리고 resnet 모델을 바꿔 가면서도 성능을 비교하는데 마찬가지로 SimCLR, MoCov2보다 더 좋은 성능을 보임을 확인할 수 있습니다.

Table 2에 대한 실험은 일부분의 labeled training set을 가지고 feature extractor를 fine-tuning 시키는(freeze X) semi-supervised training 실험을 수행하였습니다. 이때 labeled training set는 1% , 10% 만을 사용합니다. 결과 표를 보면 SimCLR 보다 좋은 성능을 보여주고 있고 Supervised 학습 방식보다 높은 성능을 보입니다.

다음 실험으로는 ImageNet 데이터셋으로 학습 시킨 모델을 다른 데이터셋으로 Transfer 시키는 실험입니다. 다른 classification 데이터셋으론 Food101, CIFAR 등 12가지 데이터셋을 사용하였고, linear evaluation과 fine-tuning 실험 결과를 보여주고 있습니다.

위의 그림의 Table 3을 보면 대부분의 경우에서 BYOL이 좋은 성능을 보여주고 있으며 Fine-tuning 실험에서 Supervised-IN의 결과보다 BYOL의 결과가 더 좋은 것을 확인할 수 있습니다.

마지막으로, Object Detection, Semantic Segmentation, Depth Estimation 등 다른 task에 feature extractor로 사용하였을 때의 실험도 진행하는데 Table 4를 보면 ImageNet으로 Pre-training을 시켰을 때보다 BYOL을 사용하였을 때 더 좋은 성능을 보이고, SimCLR 보다 더 좋은 성능을 보이는 것을 확인할 수 있습니다.

Ablation – Batch Size, Augmentations

논문에서는 Ablation study도  진행을 하는데 이 실험에서는 Contrastive Learning 방법 중 하나인 SimCLR와 비교를 하기 때문에 성능 차이가 더 두드러 지는모습을 보입니다. Fig 3 (a)에서 Batch Size가 줄어듦에 따라 급격한 성능 저하를 보이는데요. BYOL은 상대적으로 Batch Size에 둔감한 모습을 보입니다. 그리고 앞서 SimCLR는 Augmentation의 조합에 따라 성능 차이가 많이 난다고 설명드렸었는데 보시는 것처럼 Crop만 사용하는 경우 성능이 급격하게 떨어지는 것을 확인할 수 있고 그에반해 BYOL은 상대적으로 Augmentation에 강인한 모습을 보입니다.

Bootstrapping

 다음은 Bootstrapping 실험 결과인데 왼쪽 그림의 경우 Target Network를 변화시키는 방법에 따른 성능 변화를 측정한 것입니다.

τ값이 1인 경우 Constant Random Network를 의미하며 이는 Target Network를 업데이트하지 않습니다. 실험 결과를 보면 Moving Average의 정도가 적당할 때 (0.99) 성능이 가장 좋은 모습을 볼 수 있습니다.

오른쪽 table 은 contrastive loss, predictor 및 target network 의 유무에 따른 ablation 을 나타내고 있습니다. 

β=1 인 경우가 일반적인 InfoNCE loss 이고, β=0 이 BYOL 의 loss function 이라고 생각하면 됩니다. 이 결과를 통해서 BYOL의 predictor 및 target network의 유무가 매우 중요하다는 것을 확인할 수 있네요.