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

논문: https://proceedings.mlr.press/v119/chen20j.html

저자: Ting Chen, Simon Kornblith, Mohammad Norouzi, Geoffrey Hinton (Google Research, Brain Team)

인용: Chen, Ting, et al. "A simple framework for contrastive learning of visual representations." International conference on machine learning. PmLR, 2020.

깃허브: https://github.com/google-research/simclr

0. 초록 (Abstract)

 저자는 최근에 대두된 대조적 자기 지도 학습 (contrastive self-supervised learning)을 소개한다. 이는 특별한 아키텍처나 메모리 뱅크가 필요하지 않다. 이 방법이 유용한지를 살펴보기 위해 저자는 아래 세 가지를 연구 내용을 소개한다.

1. 효과적인 역할을 할 데이터 증강 모음
2. Representation ($h=t(x)$)과 contrastive loss ($z=g(h)$) 사이의 비선형 변환 (nonlinear trasnformation)의 중요성
3. 더 큰 배치 사이즈와 더 많은 훈련 기간의 영향

위 세 가지의 조합을 통해 저자는 ImageNet 데이터를 사용한 자기지도학습과 준지도학습에서 높은 성능을 달성했다. (Figure 1.)

 

1. 소개 (Introduction)

모델이 사람의 개입 (labeling)없이 효과적으로 시각적 표현 (visual representation)을 학습하는 것은 중요한 문제이다. 이에 대한 대부분의 접근법은 (1) generative 혹은 (2) discriminative이다. 전자인 생성 모델은 VAE, GAN 등으로 이미지를 생성한다. 하지만 이는 계산 비용이 높고 시각적 특징 (표현)을 학습하는데 필수적이지 않다. 후자인 분별 학습은 비지도 학습에서 레이블이 없는 데이터로부터 레이블을 만들어서, 지도학습처럼 훈련을 진행한다! 예시로 (1) Patch 위치 예측, (2) Colorization, (3) Jigsaw Puzzle, (4) Rotation 예측 등이 있으나 이들은 일반화 (generalization) 가 낮다. 이에 저자는 새로운 비지도 대조 학습 방법인 SimCLR: a simple framework for contrastive learning of visual representation 을 소개한다.

 

2. 제안 방법 (Method)

2.1. The Contrastive Learning Framework

SimCLR은 하나의 샘플로부터 다른 방식으로 증강(augmentation)된 이미지들의 embeddding vector를 서로 당기는 contrastive loss로 표현을 학습한다. (Figure 2.)

SimCLR은 아래 네 가지 주요 요소들로 이루어진다.

• 확률적 데이터 증강 모듈은 하나의 샘플에서 두 가지 데이터 증강을 확률적(e.g. 30%)으로 적용하여 positive pair를 만든다. 예를 들어, (1) random crop&resizing, (2) random color distortion, (3) random gaussian blur이다.
• 기본 신경망 인코더 $f(\cdot)$ 는 증강된 데이터들로부터 representation vector를 만든다. 저자는 ResNet을 사용하였다.
• 소규모 신경망 변환(projection head) $g(\cdot)$ 는 representation을 contrastive loss에 적용할 space로 옮긴다. 기본 MLP구조로 $g(h_i)=W^{(2)}\sigma(W^{(1)} h_i )$ 이며 $\sigma$ 는 ReLU 비선형 활성화함수이다.
• Contrastive loss 함수는 대조 예측 과정을 수행한다. positive pair 즉, 같은 샘플에서 증강된 이미지의 representation은 같은 label로 가정하고 나머지는 다른 label로 가정한다.

N개의 샘플이 포함된 미니배치 (증강된 이미지 개수는 2N개)에서 contrastive learning은 아래와 같다:

$$\begin{align*}     & \ell_{i,j} = - \log \frac{\exp(\text{sim}(\mathbf{z}_i, \mathbf{z}_j) / \tau)} {\sum_{k=1}^{2N} \mathbb{1}_{[k \neq i]} \exp(\text{sim}(\mathbf{z}_i, \mathbf{z}_k) / \tau)}     \tag{1} \end{align*}$$

여기서 $\mathbb{1}_{[k \neq i]} \in \{0,1\}$은 indicator function으로 $1 \ \text{iff} \ k \neq i$이다. 이를 algorithm (pseudo code)으로 표현하면 다음과 같다: (Algorithm 1.)

2.2. Training with Large Batch Size

저자는 메모리 뱅크를 사용하지 않고 배치사이즈의 크기를 256부터 8192까지 늘렸다. (역자 주: 이렇게 큰 배치를 쓰려면 어마어마한 금액의 GPU가 필요하다. 실제로 저자들도 Could TPU를 32 core부터 128 core까지 사용했다고..)

2.3. Evaluation Protocl

저자는 데이터셋으로 ImageNet ILSVRC-2012 dataset과 CIFAR-10 dataet을 사용하였다. 평가를 위해서는 전이학습 (trasnfer) 방식으로 특징 추출기는 고정 (frozen)한 상태로 linear classifier만 학습에 사용하였다. 참고: linear evaluation protocol (Zhang et al., 2016; Oord et al., 2018; Bachman et al., 2019; Kolesnikov et al., 2019)

 

3. Data Augmentaion for Contrastive Representation Learning

저자들은 기존에 사용되던 대조 표현 방식에서 복잡한 모델이 아닌 데이터 증강만으로도 충분히 효과를 볼 수 있음을 제시한다. (Figure 3.) 이를 증명하기 위해 널리 사용되는 데이터 증강 기법들을 소개하였다. (Figure 4.) 그리고 이들에 대한 효과를 체계적으로 분석하였다. (Figure 5.)

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

저자는 데이터 증강의 효과를 알아보기 위해 다음 두 가지 종류의 데이터 증강을 사용하였다.

• Spatial/geometric transformation: cropping and resizing, rotation, cutout
• Appearance transformation: color distortion (including color dropping, brightness, contrast, saturation, hue), gaussian blur, sobel filtering

개별 혹은 데이터 증강 조합의 효과는 Figure 5.에서 확인할 수 있다. 효과를 확인하기 위해 저자는 ImageNet 이미지들을 동일한 해상도로 crop&resize 과정을 진행하였다. 그 후 Figure 2.의 첫 번째 brunch에만 데이터 증강을 적용한 후, contrastive learning으로 학습한 모델에서 embedding에 linear classifier로 ImageNet에 대한 classification을 적용하여 아래와 같은 성능을 확인하였다. 저자는 개별 데이터 증강보다 데이터 증강 조합이 representation 학습에 더욱 도움이 됨을 확인했다. 그리고 가장 도움이 되는 데이터 증강 조합은 바로 random croping과 random color distortion이다. (accuracy 55.8%) 그 이유로는 단독 crop은 비슷한 색 분포를 가지기 때문에 color distortion이 일반화 학습에 도움이 됨을 밝혔다. (Figure 6.)

3.2 Contrastive learning needs stronger data augmentation than supervised learning

저자는 color augmentation의 효과를 자세히 확인하기 위해 추가 실험을 진행하였다. (Table 1.) 그 결과, 더 강한 color augmentation이 비지도학습(proposed)에 도움이 되었다.

 

4. Architectures for Encoder and Head

4.1. Unsupervised contrastive learning benefits (more) from bigger models

저자는 모델의 깊이가 증가할수록 성능이 올라감을 확인함으로서 지도학습보다 비지도학습에서 더 큰 모델일수록 좋은 성능을 보임을 밝혔다. (Figure 7.) 

4.2. A nonlinear projection ead improvesthe representation quality of the layer before it

저자는 contrastive learning을 수행하기 위해 representation을 해당 space로 mapping한다. 이때, 기본(none)과 linear, non-linear 세 가지로 실험을 진행하였고, 마지막이 가장 좋은 성능이 나옴을 확인했다. (Figure 8.) 다만 figure 8. 과 별개로 classification에 대해서는 대조 학습을 위한 손실로 인해 $z=g(h)$ 보다 $h=f(x)$가 더욱 좋다고 한다.

 

5. Loss Functions and Batch Size

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

저자는 모델 훈련에 대해 여러 반복수 (epoch)에 대한 배치크기의 효과를 확인하는 실험을 진행하였다. 그 결과, epoch가 작은 경우에는 배치크기가 클 수록 좋은 성능을 보였다. 다만 epoch가 커질수록 그 효과는 줄어들었다. (Figure 9.)

 

6. Comparision with State-of-the-art

저자는 ResNet-50에서 layer의 깊이를 3가지 (width $\times$1 $\times$2 $\times$4)를 사용하였으며 1000 epochs로 훈련하였다.

Transfer learning

저자는 12개의 실생활 이미지 데이터셋에서 linear evaluation*과 fine-tuning 두 가지 성능을 평가하였다. SimCLR에 사용된 feature extractor는 ResNet-50 ($\times$4)이다. 성능 결과는 Table 8. 에서 확인할 수 있다.

*좋은 representation은 단순 선형 분류기만으로도 좋은 성능을 내야 한다는 가정

 

8. Conclusion

이번 연구에서 저자는 이미지에서 대조 표현 학습에 대한 간단한 framework와 예시들을 보여주었다. 각 요소들에 대해 살펴보고 이들의 조합에 대한 효과도 확인하였다. 저자는 이런 발견들을 통해 이전 연구들에 비해 높은 성능을 달성할 수 있었다.