[논문 리뷰] A Simple Framework for Contrastive Learning of Visual Representations

논문 링크

Abstract

이 논문은 SimCLR을 제시한다.
SimCLR은 a simple framework for contrastive learning of visual representations이다.

간단하게 특수한 architecture이나 memory bank 없이 contrastive learning을 하는 방법을 제시하였다.

contrastive learning은 대조학습으로 간단하게 특정 데이터의 양성과 음성 sample을 만들고 양성으로는 representation을 가깝게 음성으로는 representation을 멀게 만들어서 표현을 하게 만든다.

연구한 내용으로 contrastive learning이 효과적인 representation을 학습하기 위해 중요한 요소는 다음과 같다.

• data augmentation의 구성이 결과 성능에 큰 영향을 끼친다.
• representation과 contrastive loss 사이에 학습 가능한 nonlinear transformation을 넣는 것은 학습되는 representation의 품질을 향상시킨다.
• 대조학습은 batch size와 training step이 클수록 좋은 품질을 얻는다.

이 batch size에 대한 내용은 추후 MoCo에서 보완이 된다!

이러한 내용들을 토대로 매우 좋은 성능을 얻었다!

1. Introduction

human label없이 visual representation을 학습하려는 문제는 오래된 문제이다.
대부분의 주요한 접근은 2개이다.

• generative
• discriminative

generative 접근법은 입력 공간의 픽셀을 생성하거나 모델링하는 방법을 학습하는 것이다.
그러나 pixel 단위의 생성은 너무 비싸고 representation learning에 필요하지도 않다.

discriminative 접근법은 label을 data에서 만들어낸 pretext task를 학습하는, supervised learning과 유사한 방법으로 representation을 학습한다.
그러나 이런 방법은 pretext task을 만드는 휴리스틱에 의존할 수밖에 없고 이는 representation의 일반성을 해친다.

그러나 contrastive learning을 활용한 discriminative 접근법은 좋은 성능을 보인다.

이 논문에서는 visual representation의 contrastive learning을 위한 간단한 framework SimCLR을 제시한다.

SimCLR은 memory bank도 필요없고 특별한 architecture도 필요가 없고 매우 간단하다.
그러나 매우 좋은 성능을 보였다.

이 논문에서는 좋은 contrastive learning을 위한 중요한 요소들을 연구하였으며 이는 다음과 같다.

• 여러개의 data augmentation의 조합은 contrastive prediction task에서 visual represenation을 만드는데 매우 중요하다.
  또한 unsupervised가 supervised보다 strong data augmentation에서 더 많은 이익을 얻는다.
• representation과 contrastive loss(InfoNCE) 사이의 learnable nonlinear transformation은 표현 학습의 질을 높인다.
• contrastive cross entropy loss를 활용한 학습은 정규화된 embedding과 적절히 조절된 temperature 매개변수로부터 이점을 얻는다.
• contrastive learning은 큰 batch size와 긴 training에서 이득을 얻는다.
  또한 supervised와 비슷하게 deeper and wider network에서도 이득을 얻는다.

이러한 내용들을 조합해서 self-supervised와 semi-supervised에서 SOTA의 성능을 얻었다.

2. Method

2.1 The Contrastive Learning Framework

SimCLR은 위 representation을 학습할 때 위 그림과 같이 같은 데이터의 augment된 view들 사이를 가깝게 만드는 식으로 진행한다.
위 그림에서 구성된 부분과 같이 각 부분으로 framework가 구성이 되는데

• 데이터를 바꿔서 view를 생성하는 stochastic data augmentation이 처음 구성이 된다. 위 그림에서는 $\tilde x_i$와 $\tilde x_j$로 positive pair이 구성이 된다.
  이 논문에서는 3가지의 간단한 augmentation이 구성된다.
  augmentation의 구성은 성능에 큰 영향을 미친다.

  • random cropping and resize back into original
  • random color distorition
  • random Gaussian blur

• base encoder $f(\cdot)$은 augmented data에서 representation vector을 만든다.
  이 논문에서는 간단하게 ResNet으로 구성을 하였다.
  $h_i=f(\tilde x_i)=ResNet(\tilde x_i)$이고 $h_i \in \mathbb{R}^d$는 average pooling layer 다음 값이다.

• projection head $g(\cdot)$은 contrastive loss가 적용이 되는 공간으로 representation을 매핑한다.
  이 논문에서는 1개의 hidden layer와 ReLU를 가진 MLP를 사용하였다.

• 이후 contrastive prediction을 위한 contrastive loss function이 정의가 되어있다. 간단하게 $\tilde x_i$와 $\tilde x_j$가 양성 pair일 때 $\{\tilde x_k\}_{k\neq i}$에서 $\tilde x_j$를 찾는 것이다.

이렇게 하나씩 구성이 되는데 이전에 contrastive learning에서 사용되던 구성과 똑같아서 이해하기 쉬울 것이다.

이 대 $N$개의 batch가 있다고 할 때 augmentation 쌍을 각각 만들어서 $2N$개의 데이터를 만든다. 이 때 자기의 쌍이 아닌 다른 $2(N-1)$개의 쌍은 다 음성 sample이다.
그리고 sim 함수는 cosine similarity일 때
positive pair(i,j)의 loss는 다음과 같이 구성이 된다.

$\ell_{i,j} = -\log \frac{\exp(\text{sim}(z_i, z_j)/\tau)}{\sum_{k=1}^{2N} \mathbb{1}_{[k \neq i]} \exp(\text{sim}(z_i, z_k)/\tau)},$

그냥 이전 논문리뷰에서 몇번 나왔던 InfoNCE loss와 동일하다.
분자의 같은 positive pair은 가깝게 분모의 negative pair와는 멀리 만든다.

final loss는 (i,j)와 (j,i)를 포함하는 모든 positive pair의 loss 합이다.
이 loss를 normalized temperature-scaled cross entropy loss라고 한다.

이 모든 과정을 의사코드로 표현하면 다음과 같다.

2.2 Training with Large Batch Size

simple을 유지하기 위해 memory bank를 사용하지 않고 batch size를 바꿔가면서 실험을 진행하였다.
batch size $N$은 256에서 8192까지 실험을 하였으며 이 때 data augmentation때문에 batch의 2배의 sample이 실제로 사용이 된다.

그러나 큰 batch size로 SGD/Momentum with linear lr의 학습을 진행하는 것이 불안정하다는 결과가 있다.
그렇기 때문에 안정화를 위해 LARS optimizer를 사용했다고 한다.

Global BN.

BN이 distsributed training에서 보통 mean과 variance는 local device에서 이루어진다.
그리고 positive 연산도 같은 batch 사이에서 이루어진다.
이 때 mean, variance를 계산하고 normalization을 하는 과정에서 같은 batch들 사이에는 통계치가 가까워져서 이러한 통계치를 기반으로 representation을 학습하지 않고 예측이 가능해져 positive sample이 어디에 있는지 정보가 유출될 수 있다.

이러한 local연산 때문에 이루어지는 information leak을 막기 위해서 모든 device의 global mean과 variance를 이용하는 BN을 활용하였다.
위 설명을 그림으로 표현하면 위와 같다. 편리를 위해 그냥 평균화로 진행하였다. original data기준으로 비슷한 분포로 모이기 때문에 positive sample이 어디에 있는지 예측하기 편하다.

3. Data Augmentation for Contrastive Representation Learning

Data augmentation defines predictive tasks

data augmentation이 supervised와 unsupervised learning에서 사용이 되고 있지만
보통 contrastive prediction task로 고려가 되지는 않았다.
문제 정의에서 이전의 다양한 접근 방법은 architecture를 바꾸는 것이었다.

그러나 이 논문에서는 augmentation을 통해 문제를 정의한다.
이미지의 부분과 그 부분 을 포함하게 자르거나 주변 이미지를 자르는 등으로 데이터를 나누어서 가깝게 만드는 것이다.
이러한 방법은 prediction task와 architecture를 분리시킨다.

결국 augmentation으로 prediction task를 정의할 수 있다.

3.1 Composition of data augmentation operations is crucial for learning good representations

augmentation의 구성에 따라 성능이 크게 바뀐다.

이 때 성능 비교를 위해 2개의 way를 따라 진행할 때 해상도 동일화를 위해 crop and resize는 고정으로 진행이 되지만 나머지 augmentation은 1개의 view에만 진행을 하였다.
이 구성은 성능을 저하시키지만 augmentation의 영향을 비교하는 것에는 도움을 주었다.

crop과 color distortion의 조합이 가장 좋았다. 이때 cropping만 사용하면 patch의 color 분포가 너무 비슷해서 문제가 쉬워졌다. 이는 모델에게 short cut이 되고 문제를 풀기 쉬워졌다.그러나 distortion을 넣으면 분포가 달라져서 문제를 풀기 어려워졌기 때문으로 추측한다.

3.2 Contrastive learning needs stronger data augmentation than supervised learning

이 논문에서는 color augmentation의 강도를 조절하며 성능을 테스트 하였는데 supervised와는 다르게 contrastive learning은 data augmentation이 강할수록 학습이 잘 되었다.

4. Architectures for Encoder and Head

4.1 Unsupervised contrastive learning benefits (more) from bigger models

간단하게 unsupervised contrastive learning도 param이 더 크면 좋은 성능을 내었다.
또한 위 그림에서 빨간 점은 1000epoch를 학습한 것인데 많이 학습할수록 좋은 성능을 내었다.

4.2 A nonlinear projection head improves the representation quality of the layer before it

앞서 학습의 구성에서 언급했던 projection head $g(\cdot)$의 종류에 따라 성능에 영향을 주었다.
위 사진은 학습한 projection head가 없는 encoder에 3가지 방법으로 layer를 올리는 것을 시도를 했는데
1. identity mapping
2. linear projection
3. MLP with 1 hidden unit and ReLU

identity << linear< non-linear의 크기로 성능이 좋았다.
그리고 이 결과는 projection head를 가지고 실험을 해도 비슷하다.

그러나 이때 projection head의 이전에 있는 $h$가 $g(h)$보다 더 좋은 성능을 보여주었다.(10% 가량)

이 논문의 저자는 저런 결과가 나온 이유가 loss function에 있다고 보았다.
$z=g(h)$는 data 변화에 불변하게 학습된다. (augment를 가한 양성 샘플 비슷하게 학습하니까)
그렇기 때문에 g는 color나 object의 위치등 downstream task에 필요한 정보를 지울 수도 있다.
이때 non-linear을 사용함으로써 더 많은 정보가 유지가될 수 있다.
이 가설을 테스트하기 위해 h또는 MLP를 사용한 $g(h)$를 사용해서 어떤 변환이 적용이 되었는지 예측하는 과정을 진행하였다.
table의 결과를 보면 $g(h)$는 정보의 손실이 일어나고 확실하게 $h$가 더 많은 정보를 포함하고 있다.

5. Loss Functions and Batch Size

5.1 Normalized cross entropy loss with adjustable temperature works better than alternatives

InfoNCE loss를 다른 contrastive loss들과 비교한 실험
결론부터 말하자면
NT-Xent loss가 가장 좋다.
이때 공정한 비교를 위해서 수정을 가했는데 이렇게 구성된 loss 함수에서 cross entropy loss 기반의 NT-Xent를 보면 분모에 negative sample을 다 더한다. 그렇기 때문에 negative sample의 상대적인 위치가 고려되어서 가중치를 부여하는데
나머지는 그냥 1대1 비교라서 상대적인 위치가 고려되지 않는다.

그렇기 때문에 공정한 비교를 위해 semi-hard negative mining을 넣었는데 이는 모든 데이터에 대해서 구하는 것이 아니라 특정한 loss margin 내에 있고 거리가 가깝지만 positive보다는 먼 데이터에 대해서만 gradient를 구하는 것이다.

이것 외에도 dot product를 진행하는 vector에 $l_2$norm을 적용하고 hyper param tuning을 진행하는등 최대한 성능을 끌어올려도 NT-xent가 가장 좋았다.

• 또한 $l_2$ norm과 적절한 temperature scaling이 없으면 성능이 많이 떨어졌다. (대략 7%)
  

5.2. Contrastive learning benefits (more) from larger batch sizes and longer training

epoch와 batch size는 클수록 좋다.
특히 batchsize는 positive와 비교되는 negative sample의 숫자가 더 많아지기 때문에 클수록 좋다.

6. Comparison with State-of-the-art

매우 좋은 성능을 보였다.

또한 적은 label로 학습하였을 때도 매우 좋은 성능을 보였음