Introduction
자연어처리 연구에서 Transformer를 사용한 Self-supervised learning (SSL)은 표준이 되어가고 있습니다. 사이즈가 큰 Transformer를 대규모의 언어 데이터셋에 학습하면 일반적으로 더 강력한 일반화 능력을 가지게 되며 downstream task에서 향상된 성능을 보여줍니다. 하지만 컴퓨터비전 연구에서는 아직 SSL은 CNN 위주였습니다. 최근 레이블 갯수가 적은 상황에서 fully-supervised pre-training 방법을 뛰어넘는 CNN SSL 방법이 나오고 있습니다. 성공의 핵심은 같은 이미지를 서로 다르게 augment한 샘플들에 대한 representation 사이의 agreement를 최대화하는 contrastive learning 입니다.
DINO, MoCo-v3와 같이 ViT를 사용한 contrastive learning 방법도 최근 연구되고 있습니다. 이들 방법은 이미지넷의 linear probe evaluation에서 SoTA 성능을 보여주지만 이미지 패치들의 긴 시퀀스에 적용되는 full self-attention으로 인하여 연산량이 매우 많다는 문제가 있습니다.
Proposed Method
Background
최근 supervised learning 에서 적용되는 multi-stage Transformer 구조가 연구되고 있습니다. 본 논문에서는 이를 처음으로 SSL에 적용해보았습니다. 구현은 직관적이지만, 충분한 고려 없이 monolithic 구조에서 multi-stage 구조로 변경하면 self-supervised Transformer의 특성을 잃어버리게 될 수 있습니다. 이에 대한 실험 결과는 아래에서 설명합니다.
Network Architectures: From Monolithic to Multi-stage ViT Backbone
위 그림은 일반적인 monolithic 구조의 Transformer와 본 논문에서 제안하는 multi-stage 구조를 비교한 것입니다. 각각의 스테이지는 patch merging/embedding 모듈과 Transformer with sparse self-ttention 모듈로 구성됩니다.
Patch merging 모듈은 각 스테이지마다 서로 다른 역할을 수행합니다. 첫 번째 스테이지에서는 입력 RGB 이미지를 겹치지 않는 패치로 자른 후 RGB 픽셀 값을 이어붙여(concatenate) C 차원의 feature로 프로젝션시킵니다. 나중 스테이지에서는 2x2 인접 패치를 이어붙인 4C 차원의 feature에 linear layer를 적용하여 2C 차원으로 만듭니다.
Transformer with sparse self-attention 모듈은 위에서 합쳐진 feature들의 interaction을 위해 적용됩니다.
이 두 모듈은 여러번(일반적으로 4번) 적용되어 multi-stage ViT를 구성합니다.
초기 스테이지에서는 패치의 갯수가 매우 많기 때문에 연산량을 줄이기 위해 sparse self-attention을 적용하는 효율적인 Transformer 구조가 고려되었습니다. 나이브하게는 Feature map을 겹치지 않는 local window로 나누고 self-attention을 윈도우 내에서만 적용하는 아이디어가 있겠지만 이미지 전체에 self-attention을 적용할 수 있는 장점을 잃어버리게 됩니다. 이 정확도-효율성의 trade-off를 최대한 활용하기 위해 여러 방법들이 제안되었습니다. 아래 방법들은 실험을 통해 결과를 비교하였습니다.
1) Swin Transformer - 각 local feature가 self-attention에서 서로 다른 윈도우에 포함되도록 연속된 Transformer block에서 번갈아가며 적용되는 두 개의 partitioning 방법.
2) Vision Longformer - 각 local window의 feature가 8개의 인접 window에 모두 attend.
3) Convolution vision Transformer - 인접한 windows의 feature가 convolutional projection에 사용됨.
Pre-training Tasks: Delving into View with Regions
본 논문에서는 SSL 시스템을 self-attention with no labels (DINO) 알고리즘을 바탕으로 만들었습니다. 자세한 내용은 DINO 논문을 참고 바랍니다.
1) View-level task
Augmented view set 인 (student), (teacher)에 대하여, pair 가 cross-view prediction task를 구성하기 위해 만들어집니다. View-level task는 DINO에서와 동일하게 아래와 같이 loss를 설계합니다.
2) Region-level task
위의 view-level task는 "local-to-global" 대응을 오직 coarse한 수준에서만 학습할 수 있습니다. 즉 이미지의 크고 작은 crop들이 view-level에서만 매치되며, 더 작은 크기의 부분(region)들 끼리의 대응은 학습될 수 없습니다. 이 문제를 해결하기 위하여 local feature 수준에서 이들의 대응을 학습할 수 있는 학습 방법을 설계하였습니다.
여기서 는 teacher feature 중에서 student feature의 번째 요소와 cosine similarity가 가장 높은 인덱스로 정의됩니다.
이 은 BERT의 masked language modeling을 모사한 것이라 할 수 있습니다. 이렇게 local-level pre-training을 적용하는 것은 dense 비전 문제에서 contrastive learning을 사용할 때 적용된 적 있지만, 본 논문에서 제안한 것은 negative sample이 없다는 점이 차이점입니다.
Experiment
본 실험에서는 ImageNet-1K 데이터셋에 라벨 없이 unsupervised pre-training을 수행하였습니다. 기본 학습 디테일은 DINO 논문을 따랐습니다. Adamw optimizer를 사용하였고, batch size는 512, 총 학습 epoch은 300 입니다.
Comparison with Prior Art on ImageNet
SSL 알고리즘을 평가하는 주요 지표는 ImageNet-1K 데이터셋에 대하여 linear probe를 실행하는 것입니다. Linear probe는 SSL로 학습된 frozen 모델로 이미지로부터 feature를 먼저 뽑아내고, 이 feature에 대한 supervised linear classifier를 학습하고 성능을 평가하는 것을 말합니다. 아래 표의 k-NN은 학습 데이터에 대해 마찬가지로 frozen 모델에 의해 feature를 뽑고 모두 저장한 뒤, 테스트 이미지의 feature를 학습 데이터들의 feature들과 비교해 k-NN을 수행한 뒤 voting 방식으로 클래스를 유추하는 방식입니다.
결과는 아래와 같습니다.
일반적으로 self-attention에서 더 긴 시퀀스를 사용하면 더 높은 성능을 얻을 수 있습니다. DINO와 MoCo-v3에서는 패치 사이즈를 8x8 혹은 7x7로 줄여서 시퀀스 길이를 늘렸습니다. EsViT에서는 윈도우 사이즈를 키움으로써 시퀀스 길이를 키웠습니다. 결과는 아래와 같습니다.
아래 결과는 region-level task의 효과를 검증하기 위한 실험입니다. view-level task만을 사용했을 때에 비해 일관적으로 높은 성능을 보여주었습니다.
아래 결과는 DeiT, Swin, ViL, CvT 네 가지 종류의 sparse attention을 비교한 것입니다. 비교를 위해 각 스테이지의 Transformer 갯수는 동일하게 2-2-6-2로 통일헀습니다. ViL과 CvT는 Swin Transformer의 안전한 대체재로 보입니다.
Transfer Learning
두 가지 태스크를 실험했습니다.
1) 18개의 작은 분류 데이터셋: 학습된 representation을 평가하는 좋은 방법은 다른 데이터셋에서 feature는 고정하고 linear classifier를 학습해보는 것입니다. 공정한 비교를 위해 자동 hyper-parameter 튜닝이 사용되었습니다.
아래 결과는 Swin-S 모델에서 supervised 방식과 EsViT를 비교한 것입니다.
2) MSCOCO 데이터셋에서 detection & segmentation
아래는 pre-train 데이터에 따른 downstrea 실험결과 차이입니다. ImageNet-1K 데이터셋에서 결과는 ImageNet-1K에 학습한 것이 가장 성능이 좋았지만 다른 18개의 데이터셋에서의 결과는 ImageNet-22K에 학습한 것이 가장 좋았습니다.
Qualitative Studies
1) Visualization of correspondences
동일한 이미지의 서로 다른 view에 대해서, pre-trained 모델로부터 top-layer feature (, ) 를 각각 뽑습니다. 의 각 vector마다 의 벡터 중 가장 비슷한 (cosine similarity가 높은) 것을 찾습니다. 이는 학습에서 사용한 region-level task와 동일한 방식입니다. 아래 그림은 매칭 결과를 표현한 것입니다.
2) Visualization of attention maps
아래 그림은 마지막 레이어의 self-attention을 도식화한 것입니다. 주 물체(강아지)의 특정 점을 query로 하여 attention이 일어난 지점을 붉은색 타일로 표시했습니다.
DINO는 class에 특정한 attention을 학습하여 query가 전경에 있든 배경에 있든 물체 위주로 attention이 일어났습니다. EsViT에서 view-level task만을 적용한 경우 multi-stage 구조로 인해 이러한 특성이 사라졌습니다. Region-level task를 추가하였을 때에는 서로 다른 head에서 attention map과 attention이 일어난 위치가 더 다양해졌습니다.
Conclusions
본 논문에서는 두 가지 새로운 아이디어 - multi-stage Transformer with sparse self-attention, 그리고 region-matching pre-training task를 사용하여 효율적인 Transformer 기반의 SSL 시스템을 만들었습니다. 또한 본 연구를 통해 더 큰 pre-training 데이터를 효과적으로 활용하는 방법이 SSL에서 매우 중요한 요소임을 확인했습니다.
Share article