[논문 리뷰] A Simple Framework for Contrastive Learning of Visual Representations / SimCLR / Self-supervised

Self-spuervised learning 에서 좋은 성능을 내는 Contrasive learning 이라는 개념을 소개하기 위해 ICML2020에 게재된 본 논문을 설명합니다.

 

아래 사이트에 그림으로 설명이 잘 되어 있어서, 자세한 내용은 참고바랍니다.

https://amitness.com/2020/03/illustrated-simclr/ 

 

Contrasive Learning

우선, contrasive learning 은 2개의 input을 네트워크에 주입했을 때, 이들이 similar 한 input 인지 different한 input 인지를 구별해주기 위한 학습 방법입니다.  예를 들어, 아래 그림에서는 Image는 고양이와는 similar 하고 강아지, 코끼리와는 different 합니다.

 

 

그런데, 이러한 과정을 위 왼쪽 그림처럼 supervised 방식으로 학습시키면, 직접 labeling을 해야하는 번거러움이 있기 때문에(사실 이럴거면 그냥 supervised로 학습시키고 class를 직접 분류하는게 더 좋겠죠), 이를 자동화하기 위해 self-supervised learning 방식을 이용합니다.

 

전제 조건은 이미지가 feature로 잘 represent 되어야 하고, 두 이미지가 embedding된 두 feature 간의 similarity를 계산하는 메커니즘이 필요합니다.

 

여기서 feature가 잘 represent 되어야 한다는 뜻은, classification 할 때는 동일한 class의 이미지들이 feature map에서 잘 모여있는 것도 중요하지만, 다른 class와 feature space에서 구별만 잘 되면(distance 가 멀면) 분류 성능도 높아지게 됩니다. 그런데 contrasive learning 에서는 동일한 class의 이미지들 끼리는 deformation이 있더라도 feature space 상에서 매우 가깝게 위치하기를 바란다는 뜻입니다. (즉, 앉아 있는 고양이나 누워있는 고양이나 동일한 고양이 이므로 feature space 에서 distance가 가까워야 similar 하다고 말할수 있게 됩니다.)

 

SimCLR Framework

본 논문에서는, 위의 문제를 self-supervised 방식으로 모델링하기 위해 'SimCLR' 이라는 프레임워크를 제안합니다. 

이 방식은 매우 간단한데, original 이미지를 서로 다른 data augmentation(in image space)를 통해 xi, xj를 생성하고 이를 동일한 Encoder(CNN)에 넣어서 feature representation hi, hj 를 얻고, fc layer를 통과시켜 zi, zj를 얻습니다.

그리고 xi와 xj의 최종 output 인 zi와  zj의 similarity를 최대화시키는 방향으로 학습을 진행합니다.

 

SimCLR training 과정

 

아래에서 구체적인 과정을 설명하려합니다. 우선, label이 없는 수백만개의 이미지로 training corpus가 있다고 가정합니다.

 

 

먼저 raw 이미지에서 크기가 N인 batch를 생성합니다. (그림 예시에서는 N=2, 논문에서는 N=8192)

Data augmentation은 (crop + flip + color jitter + grayscale)의 조합을 적용하는 random transformation function T를 정의하여 사용합니다.

 

batch의 각 이미지에 대해 2개의 이미지 쌍을 얻기 위해 data augmentation을 적용하여 총 4개(2*N =2*2 = 4)의 이미지를 얻습니다.

 

pair의 이미지들은 encoder(논문에서는 resnet50)를 통해 각각 이미지 representation(feature hi, hj) 얻고, hi, hj를 Project Head(fc layer)에 통과시켜서 non-linear transformation을 해서 zi, zj에 projection합니다. 

 

논문에서 말은 그럴싸하게 했는데, CNN으로 embedding한 feature를 fc layer에 통과시켜서 비선형성을 추가해줘서 더 복잡한 representation도 잘 하도록 해주는 역할인 것 같습니다. CNN만 통과한 feature로 similarity를 계산하면 동일한 class 임에도 deformation이 크면(고양이의 경우 자세나 카메라 각도가 다르면) similarity가 낮아질 수도 있으니까요.

 

최종 output은 아래 그림처럼 각 augmented 이미지들에 대한 feature embedding z입니다. 이제 이 z를 이용해서 loss를 계산하여 모델을 학습시킵니다.

 

이제 feature embedding z와 loss function을 이용하여 네트워크가 이미지 쌍의 similar 와 different 를 학습하면 됩니다.

 

반응형

 

Contrasive Loss 계산

 

feature 간의 similarity는 cosine similarity를 사용하여 계산됩니다. 즉, 아래 그림처럼 두 이미지 xi, xj 에 대한 feature embedding인 zi와 zj의 consine similarity를 계산합니다.

이렇게 cosine similarity를 사용하면 위 그림처럼 같은 이미지에서 augmentation 된 이미지들 끼리는 similarity가 높고(고양이-고양이, 코끼리-코끼리) 서로 다른 이미지에서 augmentation 된 이미지들 끼리는 similarity가 낮게(고양이-코끼리) 됩니다. 

 

SimCLR은 NT-Xent loss (Normalized Temperature-Scaled Cross-Entropy Loss) 라는 contrasive loss를 사용합니다.

 

 

먼저 batch에서 augmentation된 이미지 쌍을 하나씩 가져와서 두 이미지가 유사한 확률을 얻기위해 softmax 함수를 적용합니다. 이 softmax 계산은 두번 째 augmentation 된 고양이 이미지가 pair의 첫번 째 고양이 이미지와 가장 유사할 확률을 구하는 것입니다. 여기서 batch의 나머지 모든 이미지들은 dissimilar image(negative pair)로 샘플링됩니다. 

 

 

그런 다음 위 계산에 log를 음수로 취해서 pair에 대한 loss를 계산합니다. 다른 이미지 pair도 loss를 계산하고 동일한 pair에 대해 이미지의 위치가 바뀌는 경우에도 loss를 계산합니다. 

 

마지막으로, 크기가 N=2인 batch의 모든 pair에 대한 loss를 계산하고 평균을 취합니다. Loss를 기반으로 학습하면서 encoder, projection header representation 이 향상되어 feature space에서 유사한 이미지를 더 가깝게 위치시킵니다.

 

Downstream Tasks (Transfer learning)

 

SimCLR 모델이 contrasive learning 으로 학습이 완료되면 transfer learning에 사용할 수 있습니다. 이를 위해 base encoder에서 얻은 feature representation을 사용하여 classification, detection등의 down stream task에 사용할 수 있습니다.

 

아래 figure들은 논문에 설명된 SimCLR algorithm에 대한 내용입니다.

 

NT-Xent Loss

 

아래는 data augmentation의 예시와 evaluation 결과 표입니다. augmentation은 2개를 조합한 결과가 1개를 쓰는 경우보다 좋았고 crop + color distortion이 제일 좋은 결과를 보여줍니다.