들어가며

지난 글에서는 self-supervised learning 전반과 초기 연구 방법들의 아이디어에 대해 간략히 소개했습니다. 해당 글에서 언급한 대로, 초기 방법들은 supervised learning 기반 방법들의 성능을 뛰어넘지는 못합니다. 이 분야가 뛰어난 범용성을 갖는 모델을 목표로 하고 있지만 뛰어난 성능을 갖추는 것 또한 중요합니다. 따라서 이번 글에서는 초기 연구 이후 현재까지 성능 개선을 위해 진행되어온 SSL 연구 분야의 동향 및 아이디어를 소개하고, 해당 아이디어를 기반으로 뛰어난 성능을 보이는 몇 가지 방법들을 소개하겠습니다.

Contrastive learning

일반적으로, 같은 이미지에서 추출된 두 patch는 random하게 augmentation을 수행했다고 하더라도 근본적으로 포함하고 있는 context는 거의 유사할 가능성이 높습니다. 반대로 다른 이미지에서 추출된 두 patch의 경우는 포함된 객체의 class가 동일하더라도 많은 부분에서 다른 정보를 포함 할 가능성이 높습니다. contrastive learning의 주요 아이디어는 이 가정을 이용하여, 동일 이미지에서 추출된 patch의 feature 유사도는 높아지고, 다른 이미지에서 추출된 patch의 유사도는 낮아지도록 모델을 학습하는 것입니다. 올바른 비교를 위해 모델은 image에 포함된 모든 정보를 이해하는 능력을 학습하게 됩니다. supervised learning에 비해 이렇다 할 성능 상 이점이 없었던 self-supervised learning에 이 아이디어를 포함하는 contrastive learning이 적용되면서부터 성능이 크게 개선되기 시작합니다.

Contrastive learning을 이용한 self-supervised learning의 기본 구조는 다음과 같습니다.

입력 이미지 x에 대해 서로 다른 augmentation(랜덤하게 patch를 샘플링하는 augmentation 포함)을 적용하고, encoder Q, K에 각각 통과 시켜 feature q와 k를 계산합니다. 이 때, encoder Q와 K는 서로 같을 수도 있고, 다를 수도 있습니다.

인덱스가 같은 q와 k는 같은 이미지에서 추출된 feature이므로 둘 사이의 유사도가 높아야 하고, 인덱스가 다른 q와 k는 다른 이미지에서 추출된 feature 이므로 둘 사이의 유사도가 낮아야 하므로, 아래와 같은 형태의 InfoNCE loss를 이용해 모델을 학습합니다. (여기서 τ는 확률 분포를 부드럽게 만드는 역할의 hyper-parameter)

그런데 한 가지 문제는, 이 구조를 이용해 좋은 성능을 달성하기 위해서는 얼마나 양질의 negative sample이 선택될 수 있는지가 매우 중요합니다. 구분하기가 쉬운 negative sample만 계속 선택된다면 모델이 유의미한 특성을 학습하기 어려워지기 때문입니다. 따라서 많은 수의 negative sample을 포함시켜야 하고 이를 위해서는 큰 batch size의 적용이 불가피한데, 문제는 샘플들 간의 조합을 사용하다 보니 굉장히 많은 컴퓨팅 리소스를 필요로 하게 됩니다.

Moco v1, v2, v3

이를 해결하기 위한 대표적인 방법으로 Facebook의 MoCo 시리즈가 있습니다.

MoCo v1은 다음과 같은 방식을 통해 위 문제를 해결하고 성능을 개선합니다.

1. Encoder k로 전달되는 gradient를 끊어 학습에 필요한 리소스를 크게 줄입니다. 대신, encoder k의 parameter는 encoder q의 parameter로 momentum update를 수행합니다.

   
    

2. q와 k는 원래 한 batch 내의 x로부터 계산되는데, 계산된 k를 저장해 두었다가 다음 반복에도 사용하는 방식을 사용합니다. 다만 모델의 학습이 진행됨에 따라 너무 오래전 결과는 일관성이 떨어질 수 있으므로 queue 방식으로 오래된 데이터는 삭제하는 방법을 사용합니다.

MoCo v2는 다음에 설명될 SimCLR로부터 영감을 받아 encoder 뒤에 MLP head를 추가하고, 다양한 augmentation을 적용하여 기존 성능을 개선합니다.

MoCo v3는 NLP 분야에서 압도적인 성능을 보이고 있는 transformer의 가능성에 주목합니다. vision 분야에서는 ViT가 CNN 대비 성능상 큰 이점을 보이지 않고 있는데, 그 이유를 저자들은 최적의 학습 세팅이 아직 연구되지 않아 불안정성이 존재하기 때문으로 판단하고 이에 대한 실험을 진행합니다.

다양한 조건으로 실험해 본 결과, 확실히 ViT의 실험에서 굉장히 많은 불안정성이 관찰되었지만 최종적으로 ViT를 사용한 경우의 성능이 가장 좋다는 것을 보였습니다.

SimCLR v1, v2

Google의 SimCLR v1의 구조는 contrastive learning의 기본 형태와 동일합니다.

인코더(f) 다음 단계에 projection head(g)를 추가하여 성능이 개선됨을 보였고, 큰 batch 사이즈가 성능에 유리함을 실험으로 보이기도 했습니다.

SimCLR v2는 SimCLR v1에 비해 훨씬 큰 규모의 네트워크를 사용하고, 더 많은 negative sample을 확보하기 위해 MoCo의 memory mechanism을 적용하고, distillation 기법을 이용해 성능을 끌어올립니다.

Non-contrastive Learning

Contrastive learning은 매우 뛰어난 성능을 달성할 수 있는 방법임에는 분명하지만, 많은 negative sample이 확보되어야 한다는 문제가 여전히 존재합니다. memory bank 등의 방법을 사용해 이를 개선하더라도, batch size가 작다면 학습에 사용되는 데이터는 이미 몇 단계 전 모델의 출력이기 때문에 일관성 문제를 해결할 수 없습니다.

BYOL

구글의 Deepmind는 다음과 같은 한 가지 간단한 실험을 진행합니다.

• ImageNet data를 이용하여 모델 A 학습
  • 랜덤하게 초기화된 encoder의 weight를 freeze하고, 뒤의 projection head만 학습
  • encoder가 학습되어 있지 않으므로 당연히 낮은 성능을 보임 (1.4%의 top-1 accuracy)
• 모델 A의 출력을 ground truth로 모델 B 전체 학습
  • 18.8% top-1 accuracy로 성능이 크게 향상됨

이 실험에서 성능이 좋지 않은 모델의 representation이라고 하더라도, 그 representation을 학습하여 더 나은 representation을 얻을 수 있다는 사실을 알 수 있습니다.

이 실험 결과를 기반으로 저자는 다음과 같은 모델 구조를 제안합니다.

이전의 contrastive learning과 많은 부분이 비슷합니다. 우선 encoder Q, K와 유사하게, online, target 두 개의 네트워크를 사용합니다. input image에 대해 다른 augmenation을 수행한 뒤 두 네트워크를 통과 시키는 것도 유사합니다. 다른 점은, 기존에는 그렇게 얻어진 feature들 간의 거리를 계산했다면 이번에는 target network의 출력과 online network의 출력이 같아지도록 학습한다는 점입니다.

target network는 이전 실험의 랜덤 초기화 모델 A와 같이 “성능은 좋지 않지만 representation을 제공” 하는 역할을 수행합니다. online network는 이전 실험의 모델 B와 같이 “모델 A의 representation을 이용하여 결과적으로 더 좋은 representation을 학습”하는 역할을 수행합니다. 그렇게 학습된 online network의 parameter를 EMA(exponential moving average)로 target network에 update합니다.

결과적으로target network의 representation으로 online network의 representation을 개선하고, online network의 parameter로 target network의 성능을 개선하는 네트워크 구조입니다.

결과적으로 저자들은 이 구조를 통해 contrastive learning의 문제점을 보완하면서도 더 뛰어난 성능의 모델을 얻을 수 있었고, 이 아이디어를 기반으로 non-contrastive learning 방법들이 발전하게 됩니다.

iBOT

NLP 분야에서 뛰어난 성능을 보이고 있는 masked language modeling을 BYOL에 추가하여 성능을 개선한 방법입니다. 기존 BYOL과 동일한 구조를 가지고 있으나, backbone으로 ViT를 사용하고, patch의 일부를 masking 하여 학습합니다.

DINO v1, v2

마지막으로 현재 SOTA인 Meta의 DINO입니다.

DINO v1 역시 BYOL에서 영감을 받았기 때문에 많은 부분이 BYOL과 유사합니다.

BYOL과 비교했을 때 다른 점은 다음과 같습니다.

• 매우 큰 모델(ResNet, ViT)을 학습한 뒤, knowledge distillation을 이용해 개선된 모델을 얻습니다.
• MSE를 사용하는 BYOL과 달리 cross entropy loss를 사용합니다.
• Collapse를 방지하기 위해 centering과 sharpening을 추가했습니다.
  • centering : teacher의 출력 중, 특정 dimension이 우세하지 못하도록 normalize
  • sharpening : Centering의 반대 효과로, loss의 temperature를 조절하여 적용
  • 위 두 가지를 적절히 조절해 최적의 성능을 달성

DINO v2는 NLP 분야의 모델들은 별도의 fine-tuning 없이도 downstream task에서 뛰어난 성능을 보이고 있는데, vision 분야에선 그 정도 성능을 달성하지 못한 이유로 데이터가 적다는 사실을 지적합니다. 기존에는 이용할 수 있는 데이터가 ImageNet 정도가 최선인 상황이었기 때문에, 자동으로 데이터를 구축할 수 있는 process를 이용하여 142M개의 데이터를 수집합니다. 데이터 수집 뿐만 아니라 iBOT의 masking 방식 등 성능 및 계산 효율 개선을 위한 다양한 방법을 도입했습니다. 더 자세한 내용은 논문을 확인하시면 좋을 것 같습니다.

결론적으로 DINO v2는 fine-tuning을 하지 않고도 supervised learning를 뛰어넘는 성능을 달성합니다.

마지막으로 DINO v2로 추출된 특징들을 patch들에 대해 PCA를 수행한 뒤 가장 주요 component 3개에 대한 방향을 색깔로 표시한 시각화 자료입니다. 그림을 보면, 우선 배경과 객체의 구분이 명확하고, 동일 객체의 같은 부분에서는 같은 색이 표시되는 것을 볼 수 있습니다. 이를 통해 DINO v2가 이미지의 semantic information을 얼마나 잘 이해하고 있는지 알 수 있습니다.

마치며

오늘은 딥러닝 분야의 최신 트렌드 중 하나인 self-supervised learning에 대해 알아보았습니다. 본문에는 성능 관련된 내용을 거의 언급하지 않았지만, self-supervised learning으로 학습된 모델들은 classification task 뿐만 아니라, detection, segmentation 등 다양한 task에도 쉽게 적용이 가능할 정도로 general한 feature를 추출할 수 있습니다. self-supervised learning 기반의 학습 방법으로 general 한 LLM들이 많이 생겨났듯이, vision 분야에서도 general한 활용이 가능한 foundation model이 나올 때가 머지 않았다는 생각이 듭니다.

Reference

• He, Kaiming, et al. "Momentum contrast for unsupervised visual representation learning." Proceedings of the IEEE/CVF conference on computer vision and pattern recognition. 2020.
• Chen, Xinlei, et al. "Improved baselines with momentum contrastive learning." arXiv preprint arXiv:2003.04297 (2020).
• Chen, X., S. Xie, and K. He. "An empirical study of training self-supervised vision transformers. In 2021 IEEE." CVF International Conference on Computer Vision (ICCV).
• Chen, Ting, et al. "A simple framework for contrastive learning of visual representations." International conference on machine learning. PMLR, 2020.
• Chen, Ting, et al. "Big self-supervised models are strong semi-supervised learners." Advances in neural information processing systems 33 (2020): 22243-22255.
• Grill, Jean-Bastien, et al. "Bootstrap your own latent-a new approach to self-supervised learning." Advances in neural information processing systems 33 (2020): 21271-21284.
• Zhou, Jinghao, et al. "ibot: Image bert pre-training with online tokenizer." arXiv preprint arXiv:2111.07832 (2021).
• Caron, Mathilde, et al. "Emerging properties in self-supervised vision transformers." Proceedings of the IEEE/CVF international conference on computer vision. 2021.
• Oquab, Maxime, et al. "Dinov2: Learning robust visual features without supervision." arXiv preprint arXiv:2304.07193 (2023).