simCLR v1 논문, simCLR v2 논문, 코드
Self Supervised Learning를 여행하는 히치하이커를 위한 안내서
저는 요즘 모두의 연구소에서 Self Supervised Learning(SSL)을 공부하고 SSL에 대해 논문을 쓰는 것을 목표로 하는 SSL LAB(쓸랩)의 연구원으로도 활동하고 있습니다. 이번 논문은 제가 Self Supervised Learning을 처음 접했던 논문이었고, 이 분야에서도 연구가 더욱 활발하게 된 계기가 되었던 논문인 것 같습니다.
미리 알아두면 좋은 지식
- Contrastive Learning
: Positive pair와 Negative pair로 구성하여 Positive pair 끼리는 거리를 좁히고, Negative pair끼리는 거리를 멀리 띄워놓는 것
- infoNCE
: Contrastive Learning에 사용되는 Loss 중 하나
- Distilling the Knowledge in a Neural Network
: 지난 번 SSL 히치하이거 시리즈에서 소개한 논문입니다. Knowledge Distillation을 소개하는 논문이고 DistilBERT 역시 큰 뼈대는 같은 방식을 사용하고 있기 때문에 한 번 읽어보시는 것이 좋습니다
SimCLR v1 : A Simple Framework for Contrastive Learning of Visual Representations
사람의 관여없이 visual representation을 학습하는 것은 오랜 문제였습니다. 많은 mainstream 방법은 generative 혹은 discriminative 방법으로 풀려고 노력하는데 generative 방법은 pixel-level generation이 비싸고 representation learning에 굳이 필요한지에 대한 의문이 있고, discriminative 방법은 supervised learning 방법을 사용하지만 이것은 label이 있어야만 학습할 수 있다는 단점이 있습니다. 하지만 소개하는 논문 SimCLR v1 에서는 label이 필요없는 contrastive representation learning 방법으로 visual representation을 하여 아래 [그림 3]과 같은 성능을 내는 것을 보여줍니다.
미리보는 아키텍처
사실 simCLR v1은 위의 이미지만 이해해도 충분합니다.
왼쪽의 알고리즘과 오른쪽 그림을 매칭하면서 읽는 것이 좋습니다. 가장 먼저 $x$는 이미지입니다. $x$라는 이미지가 서로 각각 다른 augmentation인 $t ∼ T$, $t' ∼ T$를 거쳐 $\tilde{x}_i$, $\tilde{x}_j$가 됩니다. Augmentation한 이미지 $\tilde{x}_i$, $\tilde{x}_j$를 $f(·)$ 라는 것을 거쳐 어떤 representation으로 나타나게 하는데 논문에서 $f(·)$는 base encoder라고 부르며 ResNet을 주로 사용하였다고 합니다. 그 다음 representation이 projection head라고 부르는 $g(·)$인 MLP 기반의 네트워크를 지나서 각각 $z_i$, $z_j$가 됩니다.
알고리즘에서 알 수 있듯 $z_i$, $z_j$의 similarity(논문에서는 cosine similarity 사용)를 구하고 이 값을 infoNCE loss를 사용하여 loss를 계산하는 것이 전체적인 흐름입니다.
x가 서로 다른 두 개의 augmentation을 지나 $\tilde{x}_i$, $\tilde{x}_j$가 되었을 때 이 두 이미지는 같은 이미지로부터 얻었기 때문에 positive pair로 정의합니다. 하지만 또 다른 이미지인 y로부터 $\tilde{y}_i$, $\tilde{y}_j$를 얻었을 때 $\tilde{x}_i$와 $\tilde{y}_i$, $\tilde{y}_j$는 다른 이미지로부터 얻었기 때문에 negative pair입니다. 총 N개의 배치일 때 augmentation을 각각 하기 때문에 2N개의 샘플을 얻을 수 있고, 이 중에서 1개의 positive pair, 그리고 이 positive를 전체 2N개에서 제외한 2N-2개의 negative pair를 이용해서 학습한다고 생각하면 됩니다.
Data Augmentation for Contrastive Representation Learning
기존에 contrastive learning을 위한 data augmentation을 하기 위해선 네트워크 아키텍처를 바꾸는 수고스러움이 있었는데 simCLR v1에서는 [그림 4]에서 보이는 (b), (c), (d), (e), (i)와 같은 단순한 data augmentation을 이용하였습니다. 많은 augmentation 기법 중에 (b), (c), (d), (e), (i)를 실험을 통해 골랐다고 합니다.
[그림 6]은 위에서 말한 실험 결과를 표로 나타낸 것인데 2개의 augmentation 조합을 사용했을 때 accuracy(색깔 및 숫자로 표현)가 어떻게 되는지 나타냈습니다. 논문의 저자들은 [그림 6]의 실험을 해서 [그림 4]의 오른쪽에서 봤던 그림에서 이미지 $x$에 대하여 1개의 augmentation을 사용하는 것 보다 서로 다른 2가지의 augmentation을 사용하는 것이 성능이 가장 좋다는 것을 알아냈습니다.
밝은 노랑의 경우가 가장 accuracy가 높은데 color와 crop의 조합을 사용했을 때 가장 높습니다. 논문의 저자들은 이를 아래 [그림 6]으로 분석하였습니다. Crop한 이미지에 color distortion을 사용하지 않았을 때는 대부분의 patch에 대하여 pixel intensities의 분포가 굉장히 비슷합니다. [그림 7]의 왼쪽-위 4개의 패치들은 서로 굉장히 비슷하고 왼쪽-아래 4개의 패치들도 굉장히 비슷하기 때문에 위, 아래 서로 다른 이미지가 다르다는 것은 너무 쉽게 알 수 있습니다. 따라서 모델이 정답을 쉽게 맞출 수 있습니다. 하지만 오른쪽-위, 아래와 같이 color distortion을 사용했을 때는 그 분포가 확연히 달라지기 때문에 모델이 쉽게 정답을 맞출 수 없다는 것입니다.
아래의 [그림 8]에서는 supervised learning 보다 color distortion을 강하게 줄 수록 학습 성능이 올라간다는 것을 보여주고 supervised learning은 그렇지 않다는 것을 보여줍니다. 또한 supervised learning에 도움이 되지 않았던 Blur를 추가했을 때도 더 나은 성능을 갖게 되었다는 사실을 보여줍니다.
Understanding The Non-Linear Projection Head
[그림 4]에서 알 수 있든 encoder인 ResNet을 거친 후 MLP로 이루어진 projection head를 지납니다. [그림 8]을 통해서 다른 논문에서 사용되었던 linear projection head보다 projection head를 사용할 때가 더 좋다는 것을 보여주었고, 사용하지 않았을 때보다도 훨씬 더 좋다는 것을 보여줍니다.
그리고 [그림 9]를 통해서 downstream task에서 representation을 사용할 때 [그림 4]에서 보았던 encoder를 지난 후의 $h$를 사용하는 것이 좋은지 projection head를 지난 $z$가 좋은지 t-sne 시각화를 통해 확인을 했더니 projection head를 지난 결과가 더 잘 된다는 것을 보여주었습니다.
SimCLR v2 : Big Self-Supervised Models are Strong Semi-Supervised Learners
SimCLR v1이 나오고 다른 self supervised learning 모델에 순위를 빼앗긴 후 SimCLR v2가 나왔습니다.
미리보는 아키텍처
위의 [그림 10]은 SimCLR v2의 아키텍처입니다. SimCLR v1을 학습한 후 작은 양의 라벨을 가진 데이터셋을 이용하여 supervised fine-tuning합니다. 연구에서는 SimCLR v1을 크게 세 가지 방식으로 개선하는 SimCLRv2를 제안합니다.
Unlabeled 데이터를 처음 사용할 때는 unsupervised learning을 통해 general (visual) representations을 학습하기 위해 task-agnostic 방식으로 학습됩니다. General representation을 학습하는 것이 General pretraining의 목표이기 때문에 최종 task와는 무관한 데이터를 사용하였고 그래서 task-agnostic 이라는 말을 사용합니다. 그런 다음 general representations은 supervised fine-tuning을 통해 specific task에 맞게 조정됩니다. Unlabeled data 데이터가 두 번째로 사용될 때는 예측 성능을 더욱 향상시키고 더 작은 모델을 얻기 위해 task-specific 방식으로 사용된다. 이를 위해, fine-tuned된 teacher network에서 얻은 label을 사용하여 unlabeled 데이터에 대한 student network를 학습합니다.
논문에서 제한하는 방법은 세 가지 주요 단계로 요약될 수 있습니다: pretrain, fine-tune, distill
Main Contribution
1. General pretraining의 힘을 최대한 활용하기 위해 더 큰 사이즈의 ResNet을 사용하였습니다. SimCLR v1의 가장 큰 모델은 ResNet-50(4×)이었지만 SimCLR v2은 더 깊고 넓은 모델을 사용하였고, 가장 큰 모델은 3배 더 넓은 channels과 selective kernels(SK)을 가진 ResNet-152(3x+SK)을 사용함으로서 더 높은 성능 향상을 시켰습니다.
2. Projection head의 capacity를 키우기 위해 SimCLR보다 레이어를 하나 더 추가하였고, 이 projection head는 첫 번째 단계에서 pretraining 됩니다. 그리고 두 번째 단계에서 SimCLR v1처럼 downstream task 시기에 projection head를 버리지 않고 첫 번째 linear layer까지 encoder에 포함시킨 후 middle layer를 fine-tune에 사용하여 성능을 향상시켰습니다.
3. 또 다른 SSL 논문인 MoCo의 영향을 받아 negative example을 최대한 늘리기 위해 memory network를 추가하였습니다. 이는 작은 성능 향상을 가져왔습니다.
SimCLR v1 코드
해당 코드는 공식 코드이며 SimCLR v2와 함께 있어서 조금 헷갈릴 수도 있습니다. 따라서 논문에서 제안한 아키텍처를 생각하면서 따라가는 것이 좋습니다.
코드 228번째 줄이 모델을 만드는 부분입니다. SimCLR v1을 살펴봤을 때 처럼 SimCLR v1은 Encoder(ResNet), Projection Head가 전부입니다. 이 부분을 살펴보겠습니다.
[코드 line 231-239]
self.resnet_model과 self._projection_head 객체를 생성합니다. Projection Head는 157번째 줄에 클래스가 존재합니다. 157번째 줄부터 있는 Projection Head 클래스를 잠깐 보면 FLAGS.proj_head_mode에 따라서 projection head가 'none', 'linear', 'nonlinear' 이 되도록 결정할 수 있습니다. SimCLR v1을 학습시킬 때는 'nonlinear'여야하기 때문에 이 모드를 선택해야합니다. 그 다음엔 레이어의 개수입니다. 논문과 똑같이 설정하고 싶다면 projection head를 구성하는 레이어는 총 2개가 되어야합니다.
[코드 line 241-280]
250 - 259 까지는 extra data augmentation을 해주는 것을 볼 수 있습니다. [그림 7]에서 이야기했던 부분 같습니다.262 - 269 까지는 SimCLR v1에서 encoder, projection head를 거치는 것과 같고 그 아래의 코드들은 SimCLR v2에서 unlabeled 데이터로 pretraining 하거나 labeled 데이터를 가지고 finetune을 할 때 사용하는 코드입니다.
이제 모델을 만드는 부분을 봤으니 loss 부분을 확인해봅시다.
[코드 line 35-87]
SimCLR v1을 학습 시킬 때 벡터간의 similarity를 구하고 similarity를 infoNCE에 넣어 학습했습니다. 우리가 보고 있는 add_contrastive_loss는 총 네 개의 파라미터를 받고 있는데 hidden, hidden_norm, temperature, strategy입니다. hidden은 벡터이고 hidden_norm은 normalize된 벡터 temperature는 infoNCE에서 사용되고, strategy는 tpu사용과 관련있습니다.
51 - 53번째 줄까지는 normaliz 사용의 유무와 더불어 한 개로 묶여온 서로 비교할 벡터를 hidden1, hidden2로 쪼갭니다.
74 - 70에서 [그림 3]에서 봤던 알고리즘처럼 i, j의 similarity를 계산합니다.
81 - 85에서 전에 계산한 similarity를 이용해 infoNCE loss를 계산합니다.
SSL 논문들은 대부분 수도코드로 알고리즘이 설명되어있어서 이해하기 좋은 것 같습니다. SimCLR v2는 knowledge distillation 까지 적용되어 모델들이 정말 걷잡을 수 없이 빨리 발전하고 있다는 것도 느꼈습니다. 하아
모두의 연구소에서 진행하는
"함께 콘텐츠를 제작하는 콘텐츠 크리에이터 모임"
COCRE(코크리) 2기 회원으로 제작한 글입니다.
코크리란? 🐘
Reference
SimCLR v1 유튜브 영상 : PR-231: A Simple Framework for Contrastive Learning of Visual Representations
