Self-Supervised Learning - SimCLRv2 (1)

Semi-supervised learning은 일부의 labeled data와 매우 많은 unlabeled data를 함께 활용하는 것으로 대부분의 현실 상황과 맞물려 있어 중요한 문제로 다루어져 왔습니다. 이번 포스트에서 다룰 내용은 SimCLR에서 연장된 SimCLRv2로 self-supervised learning으로 뽑아낸 representation을 이용하여 semi-supervised learning을 수행합니다. 

기본적인 방법은 대량의 unlabeled data를 task specific하지 않은 task agnostic한 방법으로 self-supervised learning으로 수행한 후 적은 labeled data로 fine-tuning을 하는 것입니다. 이 방법은 NLP 분야에서 BERT를 기본으로 많이 사용되지만 컴퓨터 비젼에서는 잘 사용되지 않았습니다. 따라서, SimCLRv2는 1) unsupervised pretrain, 2) supervised fine-tune 의 semi-supervised를 위한 패러다임을 제시합니다.

먼저 self-supervised pretraining에서는 이미지가 task agnostic한 방법으로 클래스 라벨 없이 사용되어 특정 task에 구애받지 않은 일반화된 representation을 얻습니다. 이때, SimCLRv2에서는 기존 방법보다 매우 큰 (deep and wide) neural networks를 사용합니다. 이는 큰 모델이 self-supervised learning이나 이후의 fine-tuning에서 성능상으로 큰 강점이 있기 때문입니다. 특히, labeled data가 적을 수록 큰 모델의 효과를 더 본다고 합니다.

Self-supervised pretrainig을 완료한 이후에는 labeld data를 이용해 fine-tuning을 수행합니다. 또한, SimCLRv2는 fine-tuning이후에 모델을 컴팩트하게 개량하기 위해 task-specific하게 unlabeled data를 이용하여 distillation을 수행합니다. Distillation을 위해서 fine-tuning한 모델을 teacher network로 상정하고 teacher network의 pseudo-label을 이용하여 더 작은 student network를 학습합니다. 결과적으로 semi-supervised learning에서 기존 SOTA의 top 1 accuracy보다 21.6% (1% labeled data) / 8.7% (10% labaled data) 성능 향상을 얻을 수 있었습니다.

 

Method

전체 workflow는 Figure 3과 같습니다. 먼저 unlabeled data를 task-agnostic한 방법으로 큰 CNN에 학습합니다. 이때 기본적으로 SimCLR 방법이 응용되는데, SimCLR보다 더 깊은 projection head가 이용됩니다. 이후에 labeled data를 이용해 fine-tuning되며, fine-tuning된 teacher network 으로부터 task-specific한 방법으로 unlabeled data를 이용해 모델을 distillation합니다.

Self-supervised pretraining with SimCLRv2

SimCLRv2는 SimCLR과 근본적으로 같습니다만 크게 2가지 방법으로 SimCLR을 개량합니다.

SimCLR

1) 기존보다 더욱 큰 ResNet 모델을 이용합니다. SimCLR에서 사용되었던 ResNet-50 (4x)보다 더 깊은 152-layer ResNet (3x)를 사용하며 효율적인 파라미터 사용을 위해 selective kernel과 channel-wise attention을 이용합니다.
2) Projection head $g$를 더욱 깊게 잡습니다. 또한, SimCLR은 pretraining이 끝난 후 $g$를 단순히 제거하였지만 SimCLRv2에서는 $g$의 첫 번째 레이어를 인코더에 포함시켜 fine-tuning합니다. 결과적으로 3개의 MLP layer를 projection head로 사용하면서 projection head의 첫 번째 layer부터 fine-tuning을 수행합니다. 이 방법은 기존 대비 최대 14%의 성능 향상을 이루었다고 합니다.

Self-training / knowledge distillation via unlabeled examples

Fine-tuning된 teacher network로부터 distillation을 수행합니다. 실제 labeled data를 사용하지 않고 unlabeled data를 사용하며, teacher network의 출력을 정답으로해 다음 목적함수로 훈련시킵니다. $P(y|x_i) = exp(f^{task}(x_i)[y]/\tau) / \sum_{y'} exp(f^{task} (x_i)[y']/\tau)$ 로 $\tau$는 scale을 조절하기 위한 temperature parameter입니다.

Teacher network에서부터 나온 $P^T (y|x_i)$는 distillation 과정에서 고정되며 student network만 psuedo-label cross entropy인 위 목적함수로 학습됩니다. 또한, labeled example이 많이 존재할 경우 위의 목적함수에 label을 이용한 cross entropy를 추가해서 훈련시킬 수 있습니다.

Student network의 사이즈를 teacher network와 비슷하게 잡았을 경우 성능 향상을 보이고 작게 잡았을 경우 성능이 약간 떨어지지만 컴팩트한 모델을 얻을 수 있습니다.

---

다음 포스트

[Machine Learning/Unsupervised Learning] - Self-Supervised Learning - SimCLRv2 (2)