안녕하세요, 이번에 리뷰할 논문은 Google research team에서 2020년에 발표한 SimCLR이라는 self-supervised learning 방법론입니다. 인공지능의 대부 제프리 힌튼이 공저자인게 눈에 띕니다. 이 방법론은 MoCo, BYOL, Barlow twins, MAE 등 다른 방법론들과 함께 visual recognition의 대표적인 self-supervised learning 방법 중 하나로 꼽힙니다. 기존에 읽었던 논문들보다 복잡해서 읽고 이해하는데 애를 먹었네요. 실험과 논문 구성이 상당히 체계적이어서 배울게 많은 논문이라는 생각이 들었습니다. 한 번 살펴봅시다.

Background

일단 이 방법론이 무엇인지 바로 들어가기 전에, self-supervised learning이 무엇인지 간단하게 짚고 시작하겠습니다. 최근에 다양한 알고리즘이 개발되고 컴퓨팅 리소스가 발달하면서, GT와 예측값의 차이를 이용해 모델의 최적화를 진행하는 딥러닝 방법론은 현대사회에서 빼놓을 수 없는 연구분야가 되었습니다. 하지만 특정한 딥러닝 모델을 학습시킬 때 항상 우리의 발목을 잡는 것이 있습니다. 바로 양질의 labeled data가 절대적으로 부족하다는 점입니다. 앤드류 응 교수가 ‘모델 중심의 AI에서 데이터 중심의 AI’를 강조했듯, 성능 좋은 딥러닝 모델을 만들기 위해서는 대규모 양질의 데이터가 필요합니다. 최근 핫한 chatGPT도 방법론이 새롭다기보다는 매우 큰 데이터를 학습에 때려박아 좋은 성능을 달성했습니다. 하지만 큰 규모의, 양질의, 그것도 labeling된 데이터셋을 확보하는것은 cost가 너무 큽니다. minor한 field에서는 양질은 고사하고 해당 데이터셋이 있을지부터 걱정해야 하죠. 일반인에게 annotation을 맡기면 그나마 괜찮지만, labeling에 전문적인 지식이 필요한 field에서는(영상의학 전문의를 고용해야 한다거나) labeling cost가 더욱 커질 것입니다. 이는 아주 큰 문제입니다.

때문에 제한된 데이터를 가지고 최대한의 효율을 달성하기 위한 연구가 활발하게 이뤄지고 있습니다. semi-supervised learning(소량의 labeled data를 이용해 unlabeled data에 pseudolabeling), Active learning(라벨링 효율이 높은 데이터 선별)을 이용할수도 있고, 오늘 소개할 self-supervised learning을 사용할 수도 있습니다. self-supervised learning을 간단하게 말하자면, ‘데이터에서 label을 생성할 수 있는 다른 task를 이용해 model pretrain을 진행한 이후, 원래 목표하던 downstream task에 fine-tuning하는 방법으로 모델의 표현력을 향상시켜보자’라는 컨셉입니다. 이렇게 들으면 와닿지 않으실테니 고전적인 방법론 하나를 함께 보시죠

위 그림은 ICLR 2018에서 제안된 self-supervised learning 기법입니다. 컨셉은 아주 단순합니다. 그냥 이미지 데이터를 무작위로 0,90,180,270도 회전시켜서 몇 도 돌았는지 각도 맞추기 task를 진행합니다. 모델은 이 때 스스로 label(회전 각도)을 생성해 학습에 사용하므로, labeled data가 필요하지 않습니다. (스스로 라벨을 생성해 supervised learning을 진행하므로 self-supervised learning이라고 하는구나~ 정도로 생각하시면 좋을 것 같습니다) 이런 사전학습을 하는 task를 pretext task라고 하며, 대표적인 pretext task는 위의 각도맞추기, 색상을 추론하는 colorization, jigsaw puzzel 맞추기 등이 있습니다. 이런 간단한 사전학습을 통해 어느정도 데이터의 표현력을 모델이 학습하길 기대하는 것이죠. 실제로 pretrained model을 fine-tuning하면 처음부터 학습을 진행하는것보다 성능이 상당히 개선됩니다. 표현력을 미리 학습하므로 self-supervised learning을 representation learning이라고 부르기도 하며, 라벨을 스스로 생성하기에 unlabeled data를 사용할 수 있으므로 unsupervised learning으로 보는 시각도 있습니다.

contrastive learning은 self-supervised learning기법 중 하나입니다. 유사도 계산을 통해 비슷한 데이터(positive pair)끼리는 feature embedding space에서 가깝게, 다른 데이터(negative pair)를 멀리 있게 하도록 학습시키는 방법으로 pretrain합니다. negative pair가 많을수록 학습에 유리하므로, memory bank를 활용하기도 합니다. 오늘 살펴볼 SimCLR은 새로운 contrastive learning framework를 제안한 논문입니다. 이제 논문을 살펴보겠습니다.

SimCLR : A Simple Framework for Contrastive Learning of Visual Representations

Abstract

이 논문은 SimCLR(a Simple framework for contrastive learning of visual representations)를 제안합니다. 이 framework는 특수한 구조나 memory bank 없이도 당시 제안된 contrastive self-supervised learning algorithm을 단순화시켰습니다. 또한, contrastive prediction task에서 유용한 representation을 학습할 수 있게 하는 요인을 이해하기 위해 framework의 주요 구성 요소를 체계적으로 연구했습니다.

그 결과, (1) data augmentation의 구성이 효과적인 predictive task 정의하는 데 중요한 역할을 하고, (2) representation과 contrastive loss사이에 learnable 비선형 변환을 도입하면 학습된 representation의 품질이 크게 향상되며, (3) contrastive learning은 지도 학습에 비해 더 큰 배치 크기와 더 많은 training step을 통해 이점을 얻을 수 있다는 것을 보였습니다.

이러한 연구 결과를 결합하여 ImageNet의 self-supervised 및 semi-supervised learning에 대한 이전 방법보다 훨씬 뛰어난 성능을 달성할 수 있게 되었습니다. SimCLR로 학습한 self-supervised representation에 대해 훈련된 선형 분류기는 76.5%의 top-1 accuracy를 달성했으며, 이는 이전 SOTA에 비해 7% 개선된 것으로 supervision를 받은 ResNet-50의 성능과 일치합니다. 1%의 레이블에 대해서만 fine-tuning할 경우 85.8%의 top-5위 accuracy를 달성하여 100배 적은 레이블로 AlexNet보다 뛰어난 성능을 달성했습니다.

Introduction

사람의 지도(human supervision)없이 효과적인 시각적 표현(visual representation)을 학습하는 것은 오랜 문제였고, 대부분의 주류 접근법은 1. 생성적(generative) 접근 방식과 2.판별적(discrimitive) 접근 방식 중 하나였습니다.

1. 생성적 접근법(Generative approaches)은 input space에서 픽셀을 생성하거나 모델링하는 방법을 학습하는 방법입니다. 하지만, 픽셀 수준 생성은 계산 cost가 높고, representation learning에는 필요하지 않을 수 있습니다.
2. 판별 접근법(Discriminative approaches)은 supervised learning에 사용되는 것과 유사한 objective function를 사용하여 표현을 학습하지만, input과 label이 모두 unlabeled dataset에서 파생되는 pretext task를 수행하도록 네트워크를 훈련시킵니다. 이러한 접근 방식은 대부분 heuristic에 의존하여 pretext task를 설계해왔는데, 이는 학습된 표현의 일반성을 제한할 수 있습니다. 최근 latent space에서의 contrastive learning에 기반한 판별 접근법이 큰 가능성을 보이며 SOTA를 달성하고 있습니다.

이 논문에서 저자들은 시각적 표현의 Contrastive Learning을 위한 간단한 프레임워크(SimCLR)를 제안했습니다. SimCLR은 이전 방법론들보다 우수한 성능을 보였을 뿐만 아니라, 간단하며, 특수한 구조나 memory bank를 필요로 하지도 않습니다.

저자들은 우수한 contrastive representation learning을 가능하게 하는게 무엇인지 이해하기 위해 framework의 주요 구성 요소를 체계적으로 연구하고 다음과 같은 사항을 보였습니다 :

• 효과적인 표현을 생성하는 contrastive prediction task를 정의하는데 있어 다양한 data augmentation 작업을 구성하는 것이 중요하다. 또한, unsupervised contrastive learning은 supervised learning보다 데이터 증강을 통해 이점을 가질 수 있다
• representation과 contrastive loss 사이 learnable nonlinear transformation을 도입하는 것은 학습된 표현의 질을 크게 향상시킨다
• contrastive cross entropy loss를 활용한 representation learning은 정규화된 임베딩과 적절하게 조정된 temperature parameter로 이점을 가질 수 있다.(temperature parameter는 contrastive loss를 구성할 때 나오는 parameter입니다. simclr의 구성요소 중 하나라고 생각하시면 될 것 같습니다)
• contrastive learning은 supervised learning에 비해 더 큰 batch size와 더 긴 훈련으로 이점을 가질 수 있다. supervised learning과 마찬가지로 contrastive learning도 네트워크의 depth, width를 늘려 성능 향상을 기대할 수 있다.

저자들은 이 결과들을 조합하여, ImageNet에서 self-supervised, semi-supervised learning에서 SOTA를 달성했습니다. 다른 natural image classification dataset에서 fine-tuning했을 때, SimCLR는 12개 dataset 중 10개에서 supervised baseline보다 비슷하거나 그 이상의 성능을 보여주었다고 합니다.

The Contrastive Learning Framework

SimCLR는 latent space의 contrastive loss를 통해 동일한 data example의 differently augmented views 간 agreement를 최대화하여 표현을 학습합니다. 이게 무슨 뜻이냐면, 동일한 이미지 하나에 서로 다른 augmentation을 적용해서 두 개의 벡터를 만든 후, 이 벡터를 feature space에서 가깝게 위치시키는 방향으로 학습을 시킨다고 받아들이면 되겠습니다(반대로 서로 다른 데이터에 대해서는 멀어지는 방향으로 학습하게 됩니다)

아래 framework 사진을 보면서 주요 요소에 대해 알아보겠습니다.

framework의 구조를 나타낸 것으로 논문에서 핵심적인 그림입니다. framework는 크게 다음의 4가지 요소로 구성됩니다.

1. Data Augmentation(t∼τ, t’~τ) : 하나의 image에서 stochastic data augmentation을 통해 두개의 서로 다른 augmented image를 생성하고, 이를 positive pair로 정의합니다. 그 외의 (minibatch 내부의)나머지 augmented image들은 negative pair로 정의합니다. data augmentation은 random cropping 후 resize, random color distortion, random gaussian blur를 순차적으로 적용하는 방식으로 하는데, 실험 결과 특히 random crop과 color distortion의 조합이 좋은 성능을 달성하는데 결정적인 역할을 했다고 합니다.

2. base encoder network f(·) : augmented data examples에서 representation vector를 추출합니다. base network로는 무엇을 사용해도 되지만, 저자들은 간단히 자주 사용되는 ResNet을 실험에 사용했습니다. 위 그림의 hi, hj는 ResNet의 global average pooling 이후 output입니다.

3. projection head g(·) : 두개의 linear layer 사이에 ReLU를 넣은 구조로 구성된 작은 신경망입니다. f()에서 추출한 representation을 contrastive loss공간에 매핑시키는 역할을 하는데, 앞서 구한 hi, hj를 g(·)에 태워서 vector zi, zj를 추출하여 contrastive learning을 진행하게 됩니다.

4. contrastive loss function : contrastive prediction task를 정의하는 loss입니다. contrastive prediction task는 주어진 positive pair를 최대한 같게 만들어주려고 합니다.

학습을 진행할 때, minibatch 내부의 데이터들 중 positive pair를 제외한 데이터를 negative pair로 정의하고 학습에 사용합니다. 일반적으로 contrastive learning 방식으로 학습을 진행할 때, 1. 좋은 퀄리티를 가지며 2. 충분의 많은 양의 negative pair가 필요한데, 학습이 batch 단위로 진행되기 때문에 많은 양의 negative pair를 구성하기 위해 큰 batch size를 이용하여 학습해야 합니다. SimCLR는 기본적으로 4096의 batch size(pairing 하므로 총 8192개의 sample)를 이용하여 학습 진행하게 됩니다(엄청 큽니다 ..)

positive pair (xi, xj)에 대한 loss function은 다음과 같이 정의됩니다 : 

sim(u,v) : cosine 유사도 / τ : temperature parameter(0.1, 0.5 등등 값 사용) / 최종 loss는 mini batch에서 모든 positive pair((i,j) 및 (j,i) 둘 다)에 대해 계산합니다. 저자들은 이 loss를 T-Xent loss(the normalized temperature-scaled cross entropy loss)라고 합니다.

학습 알고리즘은 다음과 같습니다.

(위의 Figure2와 함께 보시면 이해가 수월할 것입니다)

알고리즘의 동작 과정을 살펴보면, N개(batch size)의 데이터가 들어오면 이미지 한 장당 서로 다른 augmentation(t~τ)을 태워서 2N개의 sample을 만들고, 서로 다른 augmented data 각각에 1.base encoder f() (논문에서는 ResNet50)를 태워서 representation을 얻은 뒤에, 2. projection head g() (MLP 기반의 작은 네트워크인데, 논문에서 learnable nonlinear transformation이라고도 표현합니다)를 태운 후 나온 2개의 output에 대해 3. pairwise similarity를 계산합니다. 유사도를 계산한 이후 positive pair에 대해서는 서로 가까워지게(유사도가 높아지게), negative pair에 대해서는 서로 멀어지게(유사도가 낮아지는 방향으로) 학습을 진행합니다. pretrain이 완료된 후 fine-tuning을 할 때는 projection head를 제거하고 base encoder를 거친 representation을 바로 이용합니다(Ablation experiment를 진행했을 때 제거하는게 좋았다고 하며, 저자들은 이 이유로 projection head는 data transform에 invariant하게 학습되었으므로, 물체의 방향이나 색상 정보 등 downstream task에 유용할 수 있는 정보를 제거하는 효과가 있기 때문이라는 가설을 제시합니다)

일반적으로 contrastive learning으로 pretrain을 진행할 때, 좋은 퀄리티의 충분히 많은 양의 negative pair가 있어야 훈련이 효과적으로 이루어질 수 있다고 합니다. 이를 위해서 기존 모델을은 특수한 구조를 추가하거나 memory bank를 이용하기도 했는데, SimCLR에서는 프레임워크의 단순화를 위해 batch size를 매우 크게 하는 방법을 사용합니다. batch size N을 256부터 8192까지 실험에 사용하였는데, N=8192일 때는 negative pair가 16382개까지 늘어납니다 (positive pair는 두개이죠). training resource가 엄청나게 많이 필요할 것 같습니다. 실제 실험에서는 Cloud TPU로 학습을 진행했고, 32개부터 최대 128개의 core를 사용해 분산 학습을 시켰다고 합니다. 이렇게 큰 배치단위의 학습에서는 일반적인 SGD/Momentum으로 학습하는건 학습이 불안정해 질 수 있다고 하며, 실험에 LARS optimizer라는것을 사용했다고 합니다. 처음 보는 optimizer였었는데 pytorch에도 구현되어있지 않고 그냥 대규모 배치 단위 학습을 진행할 때 안정적으로 학습하기 위해 제안된 방법론이라고 생각하시면 될 것 같습니다.

실험 대부분은 ImageNet dataset으로 진행되었으며, CIFAR-10 이외에 다양한 데이터셋에 대해서도 진행되었다고 합니다. 학습된 representation은 linear evaluation protocol로 평가하는데, 이는 학습된 network는 freeze 시키고 그 위에 linear classifier만 달아서 test accuracy를 비교하는 방법입니다.

Data Augmentation for Contrastive Representation Learning

Data Augmentation은 이 framework에서 특히 핵심적인 부분입니다. data augmentation 자체는 간단하고 우리에게 익숙하지만, 저자들은 이 간단한 방법을 체계적으로 도입함으로 기존 모델들의 복잡성을 간단히 처리했습니다. 개인적으로 논문을 읽으면서 이 부분이 인상깊었습니다. 나중에 저도 이렇게 쉽고 깔끔하게 기존 문제점들을 해결하는 방법론을 고안해보고 싶네요. 살펴 보겠습니다.

Data augmentation이 supervised learning과 unsupervised representation learning에 널리 이용되고 있지만, contrastive prediction task를 정의하는데 체계적으로 고려된 적은 없었다고 합니다. 기존의 많은 방법론들은 모델의 구조를 바꿈으로 contrastive prediction task를 정의했었습니다. 네트워크 구조의 receptive field를 제한해서 global-to-local view prediction을 구성하거나, fixed image splitting splitting을 이용해서 인접 이미지 view를 예측하는 task를 고안했었습니다. 하지만 저자들은 단순히 data augmentation만을 이용해 간단하게 이를 재연했습니다(Figure3 참고).

저자들은 매우 다양한 Data Augmentation을 고려했고, 결국 SimCLR에는 1. Random crop(with flip&resize) 2. color distortion 3.Gaussian blur가 적용되었습니다. 저자가 고려했던 Augmentation에는 cropping, resizing, horizontal flip, rotation, cutout과 같은 데이터의 공간적/기하학적 변환과, color distortion(color dropping, brightness, contrast, saturation, hue), gaussian blur, sobel filtering과 같은 외형 변환이 있습니다. 다음 그림을 참고하시면 될 것 같습니다.

각각 augmentation의 효과와 구성을 이해하기 위해서, 저자들은 각 augmentation에 대한 실험을 진행했습니다(위의 framework를 사용합니다)

Experiments

각 data augmentation의 조합을 Linear evaluation한 표 입니다. 대각선은 단일 변환을 적용한 것이고, 나머지는 2개의 augmentation을 순차적으로 조합한 것입니다. 마지막 열 Average는 행 전체의 평균입니다.

표를 보면, model이 contrastive task에서 positive pair를 거의 완벽히 identitfy 할 수 있음에도 불구하고 단일 변환이 표현을 학습하기에 부족함을 확인할 수 있습니다. Augmentation을 조합하면 contrastive prediction task가 어려워지지만 representation의 질이 상당히 개선됩니다.

colorization은 제가 기존에 잘 보지 못했던 증강 기법인데, 상당히 중요함을 확인할 수 있습니다. 단순히 cropping만 하면 결국 픽셀별 색상 값 분포는 계속 유지되는데, color distortion을 통해 이 부분을 개선할 수 있다고 합니다.

Table 1는 SimCLR와 supervised learning에서 color distortion 강도가 얼마나 성능에 영향을 미치는지 나타내는 표입니다. supervised model과 비교를 해보면 정작 supervised learning에서는 color distortion의 강도를 높여도 성능 개선이 그다지 일어나지 않습니다. 이 실험은 통해서는 ‘supervised learning보다는 contrastive learning에서 data augmentation이 더 중요하다!’는 결론을 도출해 냅니다. 왜 기존 augmentation에서 colorization이 잘 등장하지 않았는지 여기를 보면서 납득했습니다. supervised learning과 contrastive learning을 비교한 실험은 더 있습니다.

이 표는 model의 depth와 width에 대한 실험으로 supervised ResNet과 비교한 실험입니다. 초록색의 supervised와 빨간색의 contrastive에 집중하면 되는데, 모델을 키워 parameter가 증가할수록 둘 다 성능이 오르긴 하지만, 저자들의 방법론이 더 가파르게 상승합니다. 여기 표에서도 supervised보다 contrastive learning이 비교적 모델 크기를 키워 성능 향상을 기대할 수 있다는것을 보여줍니다.

다음 실험은 linear projection head에 대한 실험입니다. 보면 실험과 표가 정말정말 많이 등장하는데.. SimCLR 저자들은 최적의 framework를 제안하기 위해 정말 많은 실험을 수행했고, 어째서 SimCLR를 이렇게 구성할 수 밖에 없었는지 분석적으로 보여줍니다. 보면서 이렇게 치밀하게 실험을 하는구나 .. 하면서 감탄했습니다.

위의 framework 그림에서 projection head는 base encoder 이후에 위치하며, pretrain에서만 사용하고 fine-tuning할때는 버린다고 했습니다. representation learning에서는 중요한 영향을 미치기 때문입니다. 없을 때보다는 붙였을때가, linear transformation보단 nonlinear transformation을 하는 형식으로 구성하는것이 성능이 좋았습니다.

훈련을 어떻게 구성하는것이 최적이었는지에 관한 실험도 있습니다. 아래 표를 함께 보겠습니다.

framework 구성이 완료된 후에는 다양한 epoch와 batch size에 대한 실험이 진행되었습니다. 해당 실험에서는 길게 훈련할수록, 그리고 batch size가 충분히 커야지 좋은 성능을 달성할 수 있다는것을 보였습니다. 이는 supervised learning과는 상당히 다른 양상이라고 언급합니다(supervised learning과 비교하면 epoch과 batch size가 너무 크죠)

SOTA 방법론들과의 비교 실험입니다. 기본 모델로는 ResNe50이 사용되었습니다(1x,2x,4x는 hidden layer widths라고 합니다. Figure7에서 확인할 수 있듯 크기가 다릅니다). 학습은 1000epoch만큼 진행되었다고 합니다. 기존의 SOTA 방법론들을 찍어 누르는 성능을 보여줍니다. 논문의 하이라이트는 바로 다음 표에 있습니다.

저자들은 transfer learning 성능을 12개의 natural image dataset에 대해 linear evaluation 방법과 fine-tuning 방식으로 평가했습니다. 이 표에서는 저자들의 self-supervised ResNet50(4x)모델이 fine-tuned 됐을 때 5개의 dataset에서 supervised baseline을 능가하는 성능을 보입니다. supervised baseline은 오직 Pets, Flowers 2개에서만 우세하고, 나머지 5개 dataset에 대해서는 비슷한 성능을 보였습니다.

Conclusion

저자들은 contrastive visual representation learning을 위한 간단한 framework인 SimCLR를 제안했습니다. 이를 어떻게 구성해야 하는지 아주 자세히, 체계적으로 연구하였고 각각 요소에 대한 설계 방법에 따른 차이를 보였습니다. 이런 실험을 바탕으로 조합해서 기존의 self-supervised, semi-supervised learning보다 개선된 결과를 보였습니다.